FIR DFA: fix CFG with normal arg after lambda [KT-46825]

The fix is a bit hacky, but it's very simple. In addition, it still does
not handle the case where the receiver is a lambda function. But such
case seems to be fairly rare in practice.

compiler/fir/analysis-tests/testData/resolve/cfg/postponedLambdas.dot

@@ -24,7 +24,8 @@ digraph postponedLambdas_kt {
             color=blue
             5 [label="Enter block"];
             6 [label="Access variable R|<local>/a|"];
-            7 [label="Postponed enter to lambda"];
+            7 [label="Access variable R|<local>/b|"];
+            8 [label="Postponed enter to lambda"];
             subgraph cluster_4 {
                 color=blue
                 13 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
@@ -36,8 +37,7 @@ digraph postponedLambdas_kt {
                 }
                 17 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
             }
-            8 [label="Postponed exit from lambda"];
+            9 [label="Postponed exit from lambda"];
-            9 [label="Access variable R|<local>/b|"];
             10 [label="Function call: R|/foo|(...)"];
             11 [label="Exit block"];
         }
@@ -46,10 +46,10 @@ digraph postponedLambdas_kt {
     4 -> {5};
     5 -> {6};
     6 -> {7};
-    7 -> {13};
+    7 -> {8};
-    7 -> {8} [color=red];
+    8 -> {13};
-    7 -> {13} [style=dashed];
+    8 -> {9} [color=red];
-    8 -> {9};
+    8 -> {13} [style=dashed];
     9 -> {10};
     10 -> {11};
     11 -> {12};
@@ -57,7 +57,7 @@ digraph postponedLambdas_kt {
     14 -> {15};
     15 -> {16};
     16 -> {17};
-    17 -> {8} [color=green];
+    17 -> {9} [color=green];
     17 -> {13} [color=green style=dashed];
 
 }

compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromLibrary/callsInPlace.dot

@@ -262,22 +262,9 @@ digraph callsInPlace_kt {
             86 [label="Enter block"];
             87 [label="Postponed enter to lambda"];
             subgraph cluster_24 {
-                color=blue
-                94 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
-                subgraph cluster_25 {
-                    color=blue
-                    95 [label="Enter block"];
-                    96 [label="Const: String(test_6_1)"];
-                    97 [label="Exit block"];
-                }
-                98 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
-            }
-            88 [label="Postponed exit from lambda"];
-            89 [label="Postponed enter to lambda"];
-            subgraph cluster_26 {
                 color=blue
                 99 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
-                subgraph cluster_27 {
+                subgraph cluster_25 {
                     color=blue
                     100 [label="Enter block"];
                     101 [label="Const: String(test_6_2)"];
@@ -285,6 +272,19 @@ digraph callsInPlace_kt {
                 }
                 103 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
             }
+            88 [label="Postponed exit from lambda"];
+            89 [label="Postponed enter to lambda"];
+            subgraph cluster_26 {
+                color=blue
+                94 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+                subgraph cluster_27 {
+                    color=blue
+                    95 [label="Enter block"];
+                    96 [label="Const: String(test_6_1)"];
+                    97 [label="Exit block"];
+                }
+                98 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+            }
             90 [label="Postponed exit from lambda"];
             91 [label="Function call: R|/myRun|(...)"];
             92 [label="Exit block"];
@@ -293,13 +293,13 @@ digraph callsInPlace_kt {
     }
     85 -> {86};
     86 -> {87};
-    87 -> {94};
+    87 -> {99};
     87 -> {88} [color=red];
-    87 -> {94} [style=dashed];
+    87 -> {99} [style=dashed];
     88 -> {89};
-    89 -> {99};
+    89 -> {94};
     89 -> {90} [color=red];
-    89 -> {99} [style=dashed];
+    89 -> {94} [style=dashed];
     90 -> {91};
     91 -> {92};
     92 -> {93};
@@ -307,13 +307,13 @@ digraph callsInPlace_kt {
     95 -> {96};
     96 -> {97};
     97 -> {98};
-    98 -> {88} [color=green];
+    98 -> {90} [color=green];
     98 -> {94} [color=green style=dashed];
     99 -> {103 100};
     100 -> {101};
     101 -> {102};
     102 -> {103};
-    103 -> {90} [color=green];
+    103 -> {88} [color=green];
     103 -> {99} [color=green style=dashed];
 
     subgraph cluster_28 {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/builder/FirArgumentListBuilder.kt

@@ -5,18 +5,14 @@
 
 package org.jetbrains.kotlin.fir.expressions.builder
 
-import kotlin.contracts.*
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
 import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirArgumentListImpl
-import org.jetbrains.kotlin.fir.visitors.*
+import kotlin.contracts.ExperimentalContracts
-
+import kotlin.contracts.InvocationKind
-/*
+import kotlin.contracts.contract
- * This file was generated automatically
- * DO NOT MODIFY IT MANUALLY
- */
 
 @FirBuilderDsl
 class FirArgumentListBuilder {
@@ -35,7 +31,7 @@ class FirArgumentListBuilder {
 @OptIn(ExperimentalContracts::class)
 inline fun buildArgumentList(init: FirArgumentListBuilder.() -> Unit = {}): FirArgumentList {
     contract {
-        callsInPlace(init, kotlin.contracts.InvocationKind.EXACTLY_ONCE)
+        callsInPlace(init, InvocationKind.EXACTLY_ONCE)
     }
     return FirArgumentListBuilder().apply(init).build()
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirArgumentListImpl.kt

@@ -5,15 +5,13 @@
 
 package org.jetbrains.kotlin.fir.expressions.impl
 
+import org.jetbrains.kotlin.fir.FirPureAbstractElement
 import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.visitors.*
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
-
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
-/*
- * This file was generated automatically
- * DO NOT MODIFY IT MANUALLY
- */
 
 internal class FirArgumentListImpl(
     override val source: FirSourceElement?,
@@ -29,7 +27,22 @@ internal class FirArgumentListImpl(
     }
 
     override fun <D> transformArguments(transformer: FirTransformer<D>, data: D): FirArgumentListImpl {
-        arguments.transformInplace(transformer, data)
+        // Transform all normal arguments first and then lambda to make CFG correct. See KT-46825
+        val postponedFunctionArgs = mutableListOf<Pair<Int, FirAnonymousFunction>>()
+        val iterator = arguments.listIterator()
+        while (iterator.hasNext()) {
+            val index = iterator.nextIndex()
+            val next = iterator.next() as FirPureAbstractElement
+            if (next is FirAnonymousFunction) {
+                postponedFunctionArgs += (index to next)
+                continue
+            }
+            val result = next.transform<FirExpression, D>(transformer, data)
+            iterator.set(result)
+        }
+        for ((index, lambda) in postponedFunctionArgs) {
+            arguments[index] = lambda.transform(transformer, data)
+        }
         return this
     }
 }

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/ImplementationConfigurator.kt

@@ -491,6 +491,8 @@ object ImplementationConfigurator : AbstractFirTreeImplementationConfigurator()
         }
 
         noImpl(userTypeRef)
+
+        noImpl(argumentList)
     }
 
     private fun configureAllImplementations() {

compiler/testData/diagnostics/tests/smartCasts/variables/lambdaBetweenArguments.fir.kt

@@ -5,5 +5,5 @@ fun foo(x: Int, f: () -> Unit, y: Int) {}
 fun bar() {
     var x: Int?
     x = 4
-    foo(x, { x = null; x<!UNSAFE_CALL!>.<!>hashCode() }, <!SMARTCAST_IMPOSSIBLE!>x<!>)
+    foo(x, { x = null; x<!UNSAFE_CALL!>.<!>hashCode() }, x)
 }

compiler/testData/diagnostics/testsWithStdLib/smartcasts/lambdaInCallArgs.fir.kt

@@ -32,7 +32,7 @@ fun test_3() {
         foo(
             x.length, // stable smartcast
             { x = "" },
-            <!SMARTCAST_IMPOSSIBLE!>x<!>.length  // stable smartcast
+            x.length  // stable smartcast
         )
     }
 }