FIR CFG: move the argument reordering hack to the expression transformer

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirCallResolver.kt

@@ -77,7 +77,7 @@ class FirCallResolver(
         val functionCall = if (needTransformArguments) {
             functionCall.transformExplicitReceiver().also {
                 components.dataFlowAnalyzer.enterQualifiedAccessExpression()
-                functionCall.argumentList.transformArguments(transformer, ResolutionMode.ContextDependent)
+                functionCall.replaceArgumentList(functionCall.argumentList.transform(transformer, ResolutionMode.ContextDependent))
             }
         } else {
             functionCall
@@ -539,7 +539,7 @@ class FirCallResolver(
 
     fun resolveAnnotationCall(annotation: FirAnnotationCall): FirAnnotationCall? {
         val reference = annotation.calleeReference as? FirSimpleNamedReference ?: return null
-        annotation.argumentList.transformArguments(transformer, ResolutionMode.ContextDependent)
+        annotation.replaceArgumentList(annotation.argumentList.transform(transformer, ResolutionMode.ContextDependent))
 
         val callInfo = CallInfo(
             annotation,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -662,8 +662,9 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         // In FE1.0, it's only used for the right
         // But it seems a bit inconsistent (see KT-47409)
         // Also it's kind of complicated to transform different arguments with different expectType considering current FIR structure
-        equalityOperatorCall.transformAnnotations(transformer, ResolutionMode.ContextIndependent)
-        equalityOperatorCall.argumentList.transformArguments(transformer, withExpectedType(builtinTypes.nullableAnyType))
+        equalityOperatorCall
+            .transformAnnotations(transformer, ResolutionMode.ContextIndependent)
+            .replaceArgumentList(equalityOperatorCall.argumentList.transform(transformer, withExpectedType(builtinTypes.nullableAnyType)))
         equalityOperatorCall.resultType = equalityOperatorCall.typeRef.resolvedTypeFromPrototype(builtinTypes.booleanType.type)
 
         dataFlowAnalyzer.exitEqualityOperatorCall(equalityOperatorCall)
@@ -799,8 +800,9 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         }
 
         dataFlowAnalyzer.enterCheckNotNullCall()
-        checkNotNullCall.argumentList.transformArguments(transformer, ResolutionMode.ContextDependent)
-        checkNotNullCall.transformAnnotations(transformer, ResolutionMode.ContextIndependent)
+        checkNotNullCall
+            .transformAnnotations(transformer, ResolutionMode.ContextIndependent)
+            .replaceArgumentList(checkNotNullCall.argumentList.transform(transformer, ResolutionMode.ContextDependent))
 
         var callCompleted = false
         val result = components.syntheticCallGenerator.generateCalleeForCheckNotNullCall(checkNotNullCall, resolutionContext)?.let {
@@ -1486,6 +1488,33 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         return anonymousObjectExpression
     }
 
+    private val nonLambdaArgumentTransformer = object : FirTransformer<ResolutionMode>() {
+        override fun <E : FirElement> transformElement(element: E, data: ResolutionMode): E =
+            element.transform(this@FirExpressionsResolveTransformer, data)
+
+        override fun transformAnonymousFunctionExpression(
+            anonymousFunctionExpression: FirAnonymousFunctionExpression,
+            data: ResolutionMode
+        ): FirStatement = anonymousFunctionExpression
+    }
+
+    private val lambdaArgumentTransformer = object : FirTransformer<ResolutionMode>() {
+        override fun <E : FirElement> transformElement(element: E, data: ResolutionMode): E =
+            element
+
+        override fun transformAnonymousFunctionExpression(
+            anonymousFunctionExpression: FirAnonymousFunctionExpression,
+            data: ResolutionMode
+        ): FirStatement = this@FirExpressionsResolveTransformer.transformAnonymousFunctionExpression(anonymousFunctionExpression, data)
+    }
+
+    override fun transformArgumentList(argumentList: FirArgumentList, data: ResolutionMode): FirArgumentList {
+        // Transform all normal arguments first and then lambda to make CFG correct. See KT-46825
+        return argumentList
+            .transformArguments(nonLambdaArgumentTransformer, data)
+            .transformArguments(lambdaArgumentTransformer, data)
+    }
+
     override fun transformAnonymousFunctionExpression(
         anonymousFunctionExpression: FirAnonymousFunctionExpression,
         data: ResolutionMode

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
+import org.jetbrains.kotlin.fir.visitors.transformInplace
 
 internal class FirArgumentListImpl(
     override val source: KtSourceElement?,
@@ -27,22 +28,7 @@ internal class FirArgumentListImpl(
     }
 
     override fun <D> transformArguments(transformer: FirTransformer<D>, data: D): FirArgumentListImpl {
-        // Transform all normal arguments first and then lambda to make CFG correct. See KT-46825
-        val postponedFunctionArgs = mutableListOf<Pair<Int, FirAnonymousFunctionExpression>>()
-        val iterator = arguments.listIterator()
-        while (iterator.hasNext()) {
-            val index = iterator.nextIndex()
-            val next = iterator.next() as FirPureAbstractElement
-            if (next is FirAnonymousFunctionExpression) {
-                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)
-        }
+        arguments.transformInplace(transformer, data)
         return this
     }
 }