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/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.*
-
-/*
- * This file was generated automatically
- * DO NOT MODIFY IT MANUALLY
- */
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.InvocationKind
+import kotlin.contracts.contract
 
 @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.*
-
-/*
- * This file was generated automatically
- * DO NOT MODIFY IT MANUALLY
- */
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
 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() {