[FIR] Add forgotten replacing lambdas in arguments of try and when calls

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

@@ -7,17 +7,13 @@ package org.jetbrains.kotlin.fir.resolve.transformers
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
-import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.impl.FirBlockImpl
-import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
-import org.jetbrains.kotlin.fir.expressions.FirWrappedArgumentExpression
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.scopes.impl.withReplacedConeType
 import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.*
-import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
-import org.jetbrains.kotlin.fir.visitors.FirTransformer
-import org.jetbrains.kotlin.fir.visitors.compose
 
 internal object MapArguments : FirDefaultTransformer<Map<FirElement, FirElement>>() {
     override fun <E : FirElement> transformElement(element: E, data: Map<FirElement, FirElement>): CompositeTransformResult<E> {
@@ -37,6 +33,27 @@ internal object MapArguments : FirDefaultTransformer<Map<FirElement, FirElement>
     ): CompositeTransformResult<FirStatement> {
         return (wrappedArgumentExpression.transformChildren(this, data) as FirStatement).compose()
     }
+
+    override fun transformBlock(block: FirBlock, data: Map<FirElement, FirElement>): CompositeTransformResult<FirStatement> {
+        if (block.statements.isEmpty()) return block.compose()
+        val transformedStatement = block.statements.last().transformSingle(this@MapArguments, data)
+        when (block) {
+            is FirBlockImpl -> block.statements[block.statements.size - 1] = transformedStatement
+            is FirSingleExpressionBlock -> block.statement = transformedStatement
+        }
+        return block.compose()
+    }
+
+    override fun transformCatch(catch: FirCatch, data: Map<FirElement, FirElement>): CompositeTransformResult<FirCatch> {
+        return catch.transformBlock(this, data).compose()
+    }
+
+    override fun transformWhenBranch(
+        whenBranch: FirWhenBranch,
+        data: Map<FirElement, FirElement>,
+    ): CompositeTransformResult<FirWhenBranch> {
+        return  whenBranch.transformResult(this, data).compose()
+    }
 }
 
 internal object StoreType : FirDefaultTransformer<FirTypeRef>() {

compiler/fir/resolve/testData/resolve/problems/covariantArrayAsReceiver.txt

@@ -19,7 +19,7 @@ FILE: covariantArrayAsReceiver.kt
     public final fun foo(element: R|PsiElement|, usages: R|kotlin/Array<UsageInfo>|): R|kotlin/Unit| {
         lval adjusted: R|kotlin/collections/List<UsageInfo>?| = when () {
             (R|<local>/element| is R|KtParameter|) ->  {
-                R|<local>/usages|.R|/filterNot|<R|UsageInfo|>(<L> = filterNot@fun <implicit>.<anonymous>(): <implicit> {
+                R|<local>/usages|.R|/filterNot|<R|UsageInfo|>(<L> = filterNot@fun <anonymous>(it: R|UsageInfo|): R|kotlin/Boolean| {
                     (R|<local>/it|.R|/UsageInfo.usage| is R|KtLightMethod|)
                 }
                 )

compiler/fir/resolve/testData/resolve/stdlib/smartcasts/tryWithLambdaInside.dot

@@ -47,7 +47,7 @@ digraph tryWithLambdaInside_kt {
                         15 [label="Exit function anonymousFunction"];
                     }
                     16 [label="Postponed exit from lambda"];
-                    17 [label="Function call: R|<local>/list|.R|kotlin/collections/filter|(<L> = filter@fun <implicit>.<anonymous>(): <implicit> <kind=UNKNOWN>  {
+                    17 [label="Function call: R|<local>/list|.R|kotlin/collections/filter|(<L> = filter@fun <anonymous>(it: R|kotlin/Boolean|): R|kotlin/Boolean| <kind=UNKNOWN>  {
     R|<local>/it|
 }
 )"];
@@ -68,7 +68,7 @@ digraph tryWithLambdaInside_kt {
             24 [label="Try expression exit"];
         }
         25 [label="Jump: ^testInPlace try {
-    R|<local>/list|.R|kotlin/collections/filter|<R|kotlin/Boolean|>(<L> = filter@fun <implicit>.<anonymous>(): <implicit> <kind=UNKNOWN>  {
+    R|<local>/list|.R|kotlin/collections/filter|<R|kotlin/Boolean|>(<L> = filter@fun <anonymous>(it: R|kotlin/Boolean|): R|kotlin/Boolean| <kind=UNKNOWN>  {
         R|<local>/it|
     }
     )
@@ -119,7 +119,7 @@ finally {
                     32 [label="Access variable R|<local>/list|"];
                     33 [label="Postponed enter to lambda"];
                     34 [label="Postponed exit from lambda"];
-                    35 [label="Function call: R|<local>/list|.R|/notInPlaceFilter|(<L> = notInPlaceFilter@fun <implicit>.<anonymous>(): <implicit> {
+                    35 [label="Function call: R|<local>/list|.R|/notInPlaceFilter|(<L> = notInPlaceFilter@fun <anonymous>(it: R|kotlin/Boolean|): R|kotlin/Boolean| {
     R|<local>/it|
 }
 )"];
@@ -140,7 +140,7 @@ finally {
             42 [label="Try expression exit"];
         }
         43 [label="Jump: ^testNotInPlace try {
-    R|<local>/list|.R|/notInPlaceFilter|<R|kotlin/Boolean|>(<L> = notInPlaceFilter@fun <implicit>.<anonymous>(): <implicit> {
+    R|<local>/list|.R|/notInPlaceFilter|<R|kotlin/Boolean|>(<L> = notInPlaceFilter@fun <anonymous>(it: R|kotlin/Boolean|): R|kotlin/Boolean| {
         R|<local>/it|
     }
     )
@@ -171,4 +171,14 @@ finally {
     43 -> {44} [style=dotted];
     44 -> {45} [style=dotted];
 
+    subgraph cluster_15 {
+        color=red
+        46 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+        47 [label="Access variable R|<local>/it|"];
+        48 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+    }
+
+    46 -> {47};
+    47 -> {48};
+
 }

compiler/fir/resolve/testData/resolve/stdlib/smartcasts/tryWithLambdaInside.txt

@@ -6,7 +6,7 @@ FILE: tryWithLambdaInside.kt
     }
     public final fun testInPlace(list: R|kotlin/collections/List<kotlin/Boolean>|): R|kotlin/collections/List<kotlin/Boolean>| {
         ^testInPlace try {
-            R|<local>/list|.R|kotlin/collections/filter|<R|kotlin/Boolean|>(<L> = filter@fun <implicit>.<anonymous>(): <implicit> <kind=UNKNOWN>  {
+            R|<local>/list|.R|kotlin/collections/filter|<R|kotlin/Boolean|>(<L> = filter@fun <anonymous>(it: R|kotlin/Boolean|): R|kotlin/Boolean| <kind=UNKNOWN>  {
                 R|<local>/it|
             }
             )
@@ -17,7 +17,7 @@ FILE: tryWithLambdaInside.kt
     }
     public final fun testNotInPlace(list: R|kotlin/collections/List<kotlin/Boolean>|): R|kotlin/collections/List<kotlin/Boolean>| {
         ^testNotInPlace try {
-            R|<local>/list|.R|/notInPlaceFilter|<R|kotlin/Boolean|>(<L> = notInPlaceFilter@fun <implicit>.<anonymous>(): <implicit> {
+            R|<local>/list|.R|/notInPlaceFilter|<R|kotlin/Boolean|>(<L> = notInPlaceFilter@fun <anonymous>(it: R|kotlin/Boolean|): R|kotlin/Boolean| {
                 R|<local>/it|
             }
             )