[FIR] Git rid of creating new when and try fir nodes

It's necessary because of before change we safe original `when`
  without infered type as return statement of lambda and after that
  we add incorrect constraint system from it to outer call

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -219,11 +219,11 @@ class FirCallCompletionResultsWriterTransformer(
         val initialType = subCandidate.substitutor.substituteOrNull(typeRef.type)
         val finalType = finalSubstitutor.substituteOrNull(initialType)
 
-        val resultType = typeRef.withReplacedConeType(finalType)
+        whenExpression.resultType = typeRef.withReplacedConeType(finalType)
 
-        return whenExpression.copy(
-            resultType = resultType,
-            calleeReference = FirResolvedNamedReferenceImpl(
+        return whenExpression.transformCalleeReference(
+            StoreCalleeReference,
+            FirResolvedNamedReferenceImpl(
                 calleeReference.source,
                 calleeReference.name,
                 calleeReference.candidateSymbol
@@ -245,11 +245,10 @@ class FirCallCompletionResultsWriterTransformer(
         val initialType = subCandidate.substitutor.substituteOrNull(typeRef.type)
         val finalType = finalSubstitutor.substituteOrNull(initialType)
 
-        val resultType = typeRef.withReplacedConeType(finalType)
-
-        return tryExpression.copy(
-            resultType = resultType,
-            calleeReference = FirResolvedNamedReferenceImpl(
+        tryExpression.resultType = typeRef.withReplacedConeType(finalType)
+        return tryExpression.transformCalleeReference(
+            StoreCalleeReference,
+            FirResolvedNamedReferenceImpl(
                 calleeReference.source,
                 calleeReference.name,
                 calleeReference.candidateSymbol

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSyntheticCallGenerator.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.resolve.transformers
 
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.copy
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
@@ -19,6 +20,7 @@ import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirTryExpression
 import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirFunctionCallImpl
+import org.jetbrains.kotlin.fir.references.FirReference
 import org.jetbrains.kotlin.fir.references.impl.FirStubReference
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.calls.*
@@ -34,6 +36,9 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
 import org.jetbrains.kotlin.fir.types.impl.FirTypeProjectionWithVarianceImpl
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.compose
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.types.Variance
@@ -59,7 +64,7 @@ class FirSyntheticCallGenerator(
             SyntheticCallableId.WHEN.callableName
         ) ?: return null // TODO
 
-        return whenExpression.copy(calleeReference = reference)
+        return whenExpression.transformCalleeReference(UpdateReference, reference)
     }
 
     fun generateCalleeForTryExpression(tryExpression: FirTryExpression): FirTryExpression? {
@@ -81,7 +86,7 @@ class FirSyntheticCallGenerator(
             SyntheticCallableId.TRY.callableName
         ) ?: return null // TODO
 
-        return tryExpression.copy(calleeReference = reference)
+        return tryExpression.transformCalleeReference(UpdateReference, reference)
     }
 
     fun resolveCallableReferenceWithSyntheticOuterCall(
@@ -193,3 +198,13 @@ class FirSyntheticCallGenerator(
         }
     }
 }
+
+private object UpdateReference : FirTransformer<FirNamedReferenceWithCandidate>() {
+    override fun <E : FirElement> transformElement(element: E, data: FirNamedReferenceWithCandidate): CompositeTransformResult<E> {
+        return element.compose()
+    }
+
+    override fun transformReference(reference: FirReference, data: FirNamedReferenceWithCandidate): CompositeTransformResult<FirReference> {
+        return data.compose()
+    }
+}

compiler/fir/resolve/testData/resolve/stdlib/whenAsLambdaReturnStatement.kt

@@ -0,0 +1,21 @@
+class Module
+
+fun getModule(): Module? = null
+
+fun getInt(): Int? = null
+
+fun test_1(modules: Collection<Module>, b: Boolean) {
+    val res = modules.groupBy { module ->
+            if (b) module else module
+        }
+}
+
+fun test_2() {
+    val x = run {
+        try {
+            ""
+        } finally {
+            getInt()
+        }
+    }
+}

compiler/fir/resolve/testData/resolve/stdlib/whenAsLambdaReturnStatement.txt

@@ -0,0 +1,39 @@
+FILE: whenAsLambdaReturnStatement.kt
+    public final class Module : R|kotlin/Any| {
+        public constructor(): R|Module| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public final fun getModule(): R|Module?| {
+        ^getModule Null(null)
+    }
+    public final fun getInt(): R|kotlin/Int?| {
+        ^getInt Null(null)
+    }
+    public final fun test_1(modules: R|kotlin/collections/Collection<Module>|, b: R|kotlin/Boolean|): R|kotlin/Unit| {
+        lval res: R|kotlin/collections/Map<Module, kotlin/collections/List<Module>>| = R|<local>/modules|.R|kotlin/collections/groupBy|<R|Module|, R|Module|>(<L> = groupBy@fun <anonymous>(module: R|Module|): R|Module| <kind=UNKNOWN>  {
+            when () {
+                R|<local>/b| ->  {
+                    R|<local>/module|
+                }
+                else ->  {
+                    R|<local>/module|
+                }
+            }
+
+        }
+        )
+    }
+    public final fun test_2(): R|kotlin/Unit| {
+        lval x: R|kotlin/String| = R|kotlin/run|<R|kotlin/String|>(<L> = run@fun <anonymous>(): R|kotlin/String| <kind=EXACTLY_ONCE>  {
+            try {
+                String()
+            }
+            finally {
+                R|/getInt|()
+            }
+
+        }
+        )
+    }

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithStdlibTestGenerated.java

@@ -183,6 +183,11 @@ public class FirDiagnosticsWithStdlibTestGenerated extends AbstractFirDiagnostic
         runTest("compiler/fir/resolve/testData/resolve/stdlib/unaryOperators.kt");
     }
 
+    @TestMetadata("whenAsLambdaReturnStatement.kt")
+    public void testWhenAsLambdaReturnStatement() throws Exception {
+        runTest("compiler/fir/resolve/testData/resolve/stdlib/whenAsLambdaReturnStatement.kt");
+    }
+
     @TestMetadata("compiler/fir/resolve/testData/resolve/stdlib/callableReferences")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)