FIR CFG: do not assume all arguments to inline funs are called in place

Only values for non-noinline, non-crossinline, functional type
parameters qualify.

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerFirTestdataTestGenerated.java

@@ -1159,6 +1159,12 @@ public class DiagnosisCompilerFirTestdataTestGenerated extends AbstractDiagnosis
                 runTest("compiler/fir/analysis-tests/testData/resolve/cfg/initBlockAndInPlaceLambda.kt");
             }
 
+            @Test
+            @TestMetadata("inlineLambdasAreInplace.kt")
+            public void testInlineLambdasAreInplace() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/cfg/inlineLambdasAreInplace.kt");
+            }
+
             @Test
             @TestMetadata("innerClassInAnonymousObject.kt")
             public void testInnerClassInAnonymousObject() throws Exception {

compiler/fir/analysis-tests/legacy-fir-tests/tests-gen/org/jetbrains/kotlin/fir/LazyBodyIsNotTouchedTilContractsPhaseTestGenerated.java

@@ -1001,6 +1001,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
             runTest("compiler/fir/analysis-tests/testData/resolve/cfg/initBlockAndInPlaceLambda.kt");
         }
 
+        @TestMetadata("inlineLambdasAreInplace.kt")
+        public void testInlineLambdasAreInplace() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/cfg/inlineLambdasAreInplace.kt");
+        }
+
         @TestMetadata("innerClassInAnonymousObject.kt")
         public void testInnerClassInAnonymousObject() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/cfg/innerClassInAnonymousObject.kt");

compiler/fir/analysis-tests/testData/resolve/cfg/inlineLambdasAreInplace.fir.txt

@@ -0,0 +1,55 @@
+FILE: inlineLambdasAreInplace.kt
+    public final inline fun callsInPlaceInline(x: R|() -> kotlin/Unit|): R|kotlin/Unit| {
+        ^callsInPlaceInline R|<local>/x|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()
+    }
+    public final inline fun <T> any(x: R|T|): R|T| {
+        ^any R|<local>/x|
+    }
+    public final inline fun noinline(noinline x: R|() -> kotlin/Unit|): R|() -> kotlin/Unit| {
+        ^noinline R|<local>/x|
+    }
+    public final inline fun crossinline(crossinline x: R|() -> kotlin/Unit|): R|() -> kotlin/Unit| {
+        ^crossinline fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
+            R|<local>/x|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()
+        }
+
+    }
+    public final fun testInline(): R|kotlin/Unit| {
+        lvar x: R|kotlin/String?| = Null(null)
+        R|/callsInPlaceInline|(<L> = callsInPlaceInline@fun <anonymous>(): R|kotlin/Unit| <inline=Inline, kind=UNKNOWN>  {
+            R|<local>/x| = Null(null)
+        }
+        )
+        R|<local>/x| = String()
+        R|<local>/x|.R|kotlin/String.length|
+    }
+    public final fun testGeneric(): R|kotlin/Unit| {
+        lvar x: R|kotlin/String?| = String()
+        lval lambda: R|() -> kotlin/Unit| = R|/any|<R|() -> kotlin/Unit|>(<L> = any@fun <anonymous>(): R|kotlin/Unit| <inline=Inline>  {
+            R|<local>/x| = Null(null)
+        }
+        )
+        R|<local>/x| = String()
+        R|<local>/lambda|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()
+        R|<local>/x|.<Inapplicable(UNSTABLE_SMARTCAST): kotlin/String.length>#
+    }
+    public final fun testNoinline(): R|kotlin/Unit| {
+        lvar x: R|kotlin/String?| = String()
+        lval lambda: R|() -> kotlin/Unit| = R|/noinline|(<L> = noinline@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+            R|<local>/x| = Null(null)
+        }
+        )
+        R|<local>/x| = String()
+        R|<local>/lambda|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()
+        R|<local>/x|.<Inapplicable(UNSTABLE_SMARTCAST): kotlin/String.length>#
+    }
+    public final fun testCrossinline(): R|kotlin/Unit| {
+        lvar x: R|kotlin/String?| = String()
+        lval lambda: R|() -> kotlin/Unit| = R|/crossinline|(<L> = crossinline@fun <anonymous>(): R|kotlin/Unit| <inline=CrossInline>  {
+            R|<local>/x| = Null(null)
+        }
+        )
+        R|<local>/x| = String()
+        R|<local>/lambda|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()
+        R|<local>/x|.<Inapplicable(UNSTABLE_SMARTCAST): kotlin/String.length>#
+    }

compiler/fir/analysis-tests/testData/resolve/cfg/inlineLambdasAreInplace.kt

@@ -0,0 +1,38 @@
+inline fun callsInPlaceInline(x: () -> Unit) = x()
+
+<!NOTHING_TO_INLINE!>inline<!> fun <T> any(x: T) = x
+
+<!NOTHING_TO_INLINE!>inline<!> fun noinline(noinline x: () -> Unit) = x
+
+inline fun crossinline(crossinline x: () -> Unit) = { x() }
+
+fun testInline() {
+    var x: String? = null
+    callsInPlaceInline { x = null }
+    x = ""
+    x.length // ok
+}
+
+fun testGeneric() {
+    var x: String? = ""
+    val lambda = any { x = null }
+    x = ""
+    lambda()
+    <!SMARTCAST_IMPOSSIBLE!>x<!>.length // bad
+}
+
+fun testNoinline() {
+    var x: String? = ""
+    val lambda = noinline { x = null }
+    x = ""
+    lambda()
+    <!SMARTCAST_IMPOSSIBLE!>x<!>.length // bad
+}
+
+fun testCrossinline() {
+    var x: String? = ""
+    val lambda = crossinline { x = null }
+    x = ""
+    lambda()
+    <!SMARTCAST_IMPOSSIBLE!>x<!>.length // bad
+}

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

@@ -50,13 +50,11 @@ digraph postponedLambdas_kt {
     8 -> {9 13 10};
     8 -> {13} [style=dashed];
     9 -> {10};
-    9 -> {8} [color=green style=dashed];
     10 -> {11};
     11 -> {12};
     13 -> {14};
     14 -> {15};
     15 -> {16};
     16 -> {17};
-    17 -> {9};
 
 }

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

@@ -2,7 +2,7 @@ FILE: postponedLambdas.kt
     public final inline fun foo(vararg x: R|kotlin/Array<out kotlin/Any>|): R|kotlin/Unit| {
     }
     public final fun test(a: R|kotlin/Any|, b: R|kotlin/Any|, c: R|kotlin/Any|): R|kotlin/Unit| {
-        R|/foo|(vararg(R|<local>/a|, foo@fun <anonymous>(): R|kotlin/String| <inline=Unknown, kind=UNKNOWN>  {
+        R|/foo|(vararg(R|<local>/a|, foo@fun <anonymous>(): R|kotlin/String| <inline=Unknown>  {
             ^ String()
         }
         , R|<local>/b|))

compiler/fir/analysis-tests/testData/resolve/inference/lambdaWithoutExpectedType.fir.txt

@@ -47,13 +47,13 @@ FILE: lambdaWithoutExpectedType.kt
     public abstract interface XdIssueFolder : R|kotlin/Any| {
     }
     public final fun test(): R|kotlin/Unit| {
-        lval array: R|kotlin/Array<kotlin/Any>| = R|kotlin/arrayOf|<R|kotlin/Any|>(vararg(String(Project), arrayOf@fun <anonymous>(): R|XdProject?| <inline=Unknown, kind=UNKNOWN>  {
+        lval array: R|kotlin/Array<kotlin/Any>| = R|kotlin/arrayOf|<R|kotlin/Any|>(vararg(String(Project), arrayOf@fun <anonymous>(): R|XdProject?| <inline=Unknown>  {
             ^ Q|XdProject|.R|/filter|<R|XdProject|>(<L> = filter@fun R|FilteringContext|.<anonymous>(it: R|XdProject|): R|XdSearchingNode| <inline=NoInline>  {
                 ^ (this@R|special/anonymous|, R|<local>/it|.R|/XdProject.name|).R|/FilteringContext.eq|<R|kotlin/String|>(String(foo))
             }
             ).R|/firstOrNull|<R|XdProject|>()
         }
-        , arrayOf@fun <anonymous>(<unused var>: R|XdIssue|, <unused var>: R|XdIssueFolder|): R|kotlin/Unit| <inline=Unknown, kind=UNKNOWN>  {
+        , arrayOf@fun <anonymous>(<unused var>: R|XdIssue|, <unused var>: R|XdIssueFolder|): R|kotlin/Unit| <inline=Unknown>  {
             ^@arrayOf Unit
         }
         ))

compiler/fir/analysis-tests/testData/resolveWithStdlib/delegates/delegateTypeMismatch.fir.txt

@@ -19,7 +19,7 @@ FILE: delegateTypeMismatch.kt
             public get(): R|kotlin/Boolean|
 
         private final fun <T> property(initialValue: R|T|): R|kotlin/properties/ReadWriteProperty<A, T>| {
-            ^property Q|kotlin/properties/Delegates|.R|kotlin/properties/Delegates.vetoable|<R|T|>(R|<local>/initialValue|, <L> = vetoable@fun <anonymous>(<unused var>: R|@R|kotlin/ParameterName|(name = String(property))  kotlin/reflect/KProperty<*>|, <unused var>: R|T|, <unused var>: R|T|): R|kotlin/Boolean| <inline=CrossInline, kind=UNKNOWN>  {
+            ^property Q|kotlin/properties/Delegates|.R|kotlin/properties/Delegates.vetoable|<R|T|>(R|<local>/initialValue|, <L> = vetoable@fun <anonymous>(<unused var>: R|@R|kotlin/ParameterName|(name = String(property))  kotlin/reflect/KProperty<*>|, <unused var>: R|T|, <unused var>: R|T|): R|kotlin/Boolean| <inline=CrossInline>  {
                 ^ when () {
                     this@R|/A|.R|/A.isLocked| ->  {
                         throw R|java/lang/IllegalStateException.IllegalStateException|(String(Cannot modify readonly DescriptorRendererOptions))

compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/plusAssignWithLambdaInRhs.dot

@@ -43,14 +43,12 @@ digraph plusAssignWithLambdaInRhs_kt {
     7 -> {8 12 9} [style=dotted];
     7 -> {12} [style=dashed];
     8 -> {9} [style=dotted];
-    8 -> {7} [color=green style=dotted];
     9 -> {10} [style=dotted];
     10 -> {11} [style=dotted];
     12 -> {13} [style=dotted];
     13 -> {14} [style=dotted];
     14 -> {15} [style=dotted];
     15 -> {16} [style=dotted];
-    16 -> {8} [style=dotted];
 
     subgraph cluster_4 {
         color=red
@@ -120,7 +118,7 @@ digraph plusAssignWithLambdaInRhs_kt {
             31 [label="Postponed exit from lambda"];
             32 [label="Function call: R|/A.A|<R|T|>(...)" style="filled" fillcolor=yellow];
             33 [label="Jump: ^postpone R|/A.A|<R|T|>(<L> = A@fun <anonymous>(resolve: R|(T) -> kotlin/Unit|): R|kotlin/Unit| <inline=NoInline>  {
-    R|<local>/queue|.R|kotlin/collections/plusAssign|<R|() -> kotlin/Unit|>(A@fun <anonymous>(): R|kotlin/Unit| <inline=Inline, kind=UNKNOWN>  {
+    R|<local>/queue|.R|kotlin/collections/plusAssign|<R|() -> kotlin/Unit|>(A@fun <anonymous>(): R|kotlin/Unit| <inline=Inline>  {
         R|<local>/resolve|.R|SubstitutionOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/computation|.R|SubstitutionOverride<kotlin/Function0.invoke: R|T|>|())
     }
     )
@@ -149,7 +147,6 @@ digraph plusAssignWithLambdaInRhs_kt {
     40 -> {41 45 42};
     40 -> {45} [style=dashed];
     41 -> {42};
-    41 -> {40} [color=green style=dashed];
     42 -> {43};
     43 -> {44};
     45 -> {46};
@@ -157,6 +154,5 @@ digraph plusAssignWithLambdaInRhs_kt {
     47 -> {48};
     48 -> {49};
     49 -> {50};
-    50 -> {41};
 
 }

compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/plusAssignWithLambdaInRhs.fir.txt

@@ -1,7 +1,7 @@
 FILE: plusAssignWithLambdaInRhs.kt
     public final fun test(): R|kotlin/Unit| {
         lval list: R|kotlin/collections/MutableList<kotlin/Function1<kotlin/String, kotlin/String>>| = Null(null)!!
-        R|<local>/list|.R|kotlin/collections/plusAssign|<R|(kotlin/String) -> kotlin/String|>(fun <anonymous>(it: R|kotlin/String|): R|kotlin/String| <inline=Inline, kind=UNKNOWN>  {
+        R|<local>/list|.R|kotlin/collections/plusAssign|<R|(kotlin/String) -> kotlin/String|>(fun <anonymous>(it: R|kotlin/String|): R|kotlin/String| <inline=Inline>  {
             ^ R|<local>/it|
         }
         )
@@ -18,7 +18,7 @@ FILE: plusAssignWithLambdaInRhs.kt
     public final fun <T> postpone(computation: R|() -> T|): R|A<T>| {
         lval queue: R|kotlin/collections/MutableList<kotlin/Function0<kotlin/Unit>>| = R|kotlin/collections/mutableListOf|<R|() -> kotlin/Unit|>()
         ^postpone R|/A.A|<R|T|>(<L> = A@fun <anonymous>(resolve: R|(T) -> kotlin/Unit|): R|kotlin/Unit| <inline=NoInline>  {
-            R|<local>/queue|.R|kotlin/collections/plusAssign|<R|() -> kotlin/Unit|>(A@fun <anonymous>(): R|kotlin/Unit| <inline=Inline, kind=UNKNOWN>  {
+            R|<local>/queue|.R|kotlin/collections/plusAssign|<R|() -> kotlin/Unit|>(A@fun <anonymous>(): R|kotlin/Unit| <inline=Inline>  {
                 R|<local>/resolve|.R|SubstitutionOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/computation|.R|SubstitutionOverride<kotlin/Function0.invoke: R|T|>|())
             }
             )

compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/problems/buildListLazy.fir.txt

@@ -28,7 +28,7 @@ FILE: buildListLazy.kt
                 this@R|special/anonymous|.R|kotlin/collections/plusAssign|<R|NameAndSafeValue|>(R|/NameAndSafeValue.NameAndSafeValue|(R|<local>/name|, R|<local>/value|))
             }
             )
-            this@R|special/anonymous|.R|kotlin/collections/sortBy|<R|NameAndSafeValue|, R|kotlin/String|>(<L> = sortBy@fun <anonymous>(it: R|NameAndSafeValue|): R|kotlin/String?| <inline=CrossInline, kind=UNKNOWN>  {
+            this@R|special/anonymous|.R|kotlin/collections/sortBy|<R|NameAndSafeValue|, R|kotlin/String|>(<L> = sortBy@fun <anonymous>(it: R|NameAndSafeValue|): R|kotlin/String?| <inline=CrossInline>  {
                 ^ R|<local>/it|.R|/NameAndSafeValue.name|
             }
             )

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticTestGenerated.java

@@ -1159,6 +1159,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
                 runTest("compiler/fir/analysis-tests/testData/resolve/cfg/initBlockAndInPlaceLambda.kt");
             }
 
+            @Test
+            @TestMetadata("inlineLambdasAreInplace.kt")
+            public void testInlineLambdasAreInplace() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/cfg/inlineLambdasAreInplace.kt");
+            }
+
             @Test
             @TestMetadata("innerClassInAnonymousObject.kt")
             public void testInnerClassInAnonymousObject() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -1159,6 +1159,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
                 runTest("compiler/fir/analysis-tests/testData/resolve/cfg/initBlockAndInPlaceLambda.kt");
             }
 
+            @Test
+            @TestMetadata("inlineLambdasAreInplace.kt")
+            public void testInlineLambdasAreInplace() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/cfg/inlineLambdasAreInplace.kt");
+            }
+
             @Test
             @TestMetadata("innerClassInAnonymousObject.kt")
             public void testInnerClassInAnonymousObject() throws Exception {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutorByMap
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
+import org.jetbrains.kotlin.fir.resolve.transformers.unwrapAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.scopes.getFunctions
 import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
@@ -808,12 +809,6 @@ abstract class FirDataFlowAnalyzer(
         graphBuilder.exitResolvedQualifierNode(resolvedQualifier).mergeIncomingFlow()
     }
 
-    private tailrec fun FirExpression.getAnonymousFunction(): FirAnonymousFunction? = when (this) {
-        is FirAnonymousFunctionExpression -> anonymousFunction
-        is FirLambdaArgumentExpression -> expression.getAnonymousFunction()
-        else -> null
-    }
-
     private var functionCallLevel = 0
     private var resolvingAugmentedAssignmentOptions: Boolean = false
 
@@ -844,7 +839,7 @@ abstract class FirDataFlowAnalyzer(
     }
 
     fun enterFunctionCall(functionCall: FirFunctionCall) {
-        val lambdaArgs = functionCall.arguments.mapNotNullTo(mutableSetOf()) { it.getAnonymousFunction() }
+        val lambdaArgs = functionCall.arguments.mapNotNullTo(mutableSetOf()) { it.unwrapAnonymousFunctionExpression() }
         val localVariableAssignmentAnalyzer = context.firLocalVariableAssignmentAnalyzer
             ?: if (lambdaArgs.isNotEmpty()) getOrCreateLocalVariableAssignmentAnalyzer(lambdaArgs.first()) else null
         localVariableAssignmentAnalyzer?.enterFunctionCall(lambdaArgs, functionCallLevel)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt

@@ -11,9 +11,7 @@ import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.FirFunction
-import org.jetbrains.kotlin.fir.declarations.FirReceiverParameter
 import org.jetbrains.kotlin.fir.declarations.builder.buildContextReceiver
-import org.jetbrains.kotlin.fir.declarations.builder.buildReceiverParameterCopy
 import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
@@ -24,11 +22,10 @@ import org.jetbrains.kotlin.fir.resolve.inference.model.ConeExpectedTypeConstrai
 import org.jetbrains.kotlin.fir.resolve.initialTypeOfCandidate
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.FirCallCompletionResultsWriterTransformer
-import org.jetbrains.kotlin.fir.resolve.transformers.InvocationKindTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformerDispatcher
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
+import org.jetbrains.kotlin.fir.resolve.transformers.replaceLambdaArgumentInvocationKinds
 import org.jetbrains.kotlin.fir.resolve.typeFromCallee
 import org.jetbrains.kotlin.fir.symbols.impl.FirValueParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
@@ -45,8 +42,6 @@ import org.jetbrains.kotlin.types.model.StubTypeMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
 import org.jetbrains.kotlin.types.model.safeSubstitute
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
-import kotlin.contracts.InvocationKind
-import kotlin.contracts.contract
 
 class FirCallCompleter(
     private val transformer: FirAbstractBodyResolveTransformerDispatcher,
@@ -123,7 +118,9 @@ class FirCallCompleter(
         val completionMode = candidate.computeCompletionMode(session.inferenceComponents, expectedTypeRef, initialType)
 
         val analyzer = createPostponedArgumentsAnalyzer(transformer.resolutionContext)
-        call.transformSingle(InvocationKindTransformer, null)
+        if (call is FirFunctionCall) {
+            call.replaceLambdaArgumentInvocationKinds(session)
+        }
 
         return when (completionMode) {
             ConstraintSystemCompletionMode.FULL -> {

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

@@ -6,98 +6,49 @@
 package org.jetbrains.kotlin.fir.resolve.transformers
 
 import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
-import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.contracts.description.ConeCallsEffectDeclaration
 import org.jetbrains.kotlin.fir.contracts.effects
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.declarations.utils.isInline
-import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirWrappedArgumentExpression
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
-import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.isFunctionalOrSuspendFunctionalType
 
-object InvocationKindTransformer : FirTransformer<Any?>() {
+tailrec fun FirExpression.unwrapAnonymousFunctionExpression(): FirAnonymousFunction? = when (this) {
-    private object ArgumentsTransformer : FirTransformer<Pair<Map<FirExpression, EventOccurrencesRange>, EventOccurrencesRange?>>() {
+    is FirAnonymousFunctionExpression -> anonymousFunction
-        override fun <E : FirElement> transformElement(element: E, data: Pair<Map<FirExpression, EventOccurrencesRange>, EventOccurrencesRange?>): E {
+    is FirWrappedArgumentExpression -> expression.unwrapAnonymousFunctionExpression()
-            return element
+    else -> null
-        }
+}
 
-        override fun transformAnonymousFunctionExpression(
+fun FirFunctionCall.replaceLambdaArgumentInvocationKinds(session: FirSession) {
-            anonymousFunctionExpression: FirAnonymousFunctionExpression,
+    val calleeReference = calleeReference as? FirNamedReferenceWithCandidate ?: return
-            data: Pair<Map<FirExpression, EventOccurrencesRange>, EventOccurrencesRange?>
+    val argumentMapping = calleeReference.candidate.argumentMapping ?: return
-        ): FirStatement {
+    val function = calleeReference.candidateSymbol.fir as? FirSimpleFunction ?: return
-            val kind = data.second ?: data.first[anonymousFunctionExpression]
+    val isInline = function.isInline
-            return anonymousFunctionExpression.transformAnonymousFunction(this, emptyMap<FirExpression, EventOccurrencesRange>() to kind)
-        }
 
-        override fun transformAnonymousFunction(
+    val byParameter = mutableMapOf<FirValueParameter, EventOccurrencesRange>()
-            anonymousFunction: FirAnonymousFunction,
+    function.contractDescription.effects?.forEach { fir ->
-            data: Pair<Map<FirExpression, EventOccurrencesRange>, EventOccurrencesRange?>
+        val effect = fir.effect as? ConeCallsEffectDeclaration ?: return@forEach
-        ): FirStatement {
+        // TODO: Support callsInPlace contracts on receivers
-            val kind = data.second
+        val valueParameter = function.valueParameters.getOrNull(effect.valueParameterReference.parameterIndex) ?: return@forEach
-            if (kind != null) {
+        byParameter[valueParameter] = effect.kind
-                anonymousFunction.replaceInvocationKind(kind)
-            }
-            return anonymousFunction
-        }
-
-        override fun transformLambdaArgumentExpression(
-            lambdaArgumentExpression: FirLambdaArgumentExpression,
-            data: Pair<Map<FirExpression, EventOccurrencesRange>, EventOccurrencesRange?>
-        ): FirStatement {
-            return data.first[lambdaArgumentExpression]?.let {
-                (lambdaArgumentExpression.transformChildren(this, data.first to it) as FirStatement)
-            } ?: lambdaArgumentExpression
-        }
-
-        override fun transformNamedArgumentExpression(
-            namedArgumentExpression: FirNamedArgumentExpression,
-            data: Pair<Map<FirExpression, EventOccurrencesRange>, EventOccurrencesRange?>
-        ): FirStatement {
-            return data.first[namedArgumentExpression]?.let {
-                (namedArgumentExpression.transformChildren(this, data.first to it) as FirStatement)
-            } ?: namedArgumentExpression
-        }
     }
+    if (byParameter.isEmpty() && !isInline) return
 
-    override fun <E : FirElement> transformElement(element: E, data: Any?): E {
+    for ((argument, parameter) in argumentMapping) {
-        return element
+        val lambda = argument.unwrapAnonymousFunctionExpression() ?: continue
-    }
+        val kind = byParameter[parameter] ?: EventOccurrencesRange.UNKNOWN.takeIf {
-
+            // Inline functional parameters have to be called in-place; that's the only permitted operation on them.
-    override fun transformFunctionCall(functionCall: FirFunctionCall, data: Any?): FirStatement {
+            isInline && !parameter.isNoinline && !parameter.isCrossinline &&
-        val calleeReference = functionCall.calleeReference as? FirNamedReferenceWithCandidate ?: return functionCall
+                    parameter.returnTypeRef.coneType.isFunctionalOrSuspendFunctionalType(session)
-        val argumentMapping = calleeReference.candidate.argumentMapping ?: return functionCall
+        } ?: continue
-        val function = calleeReference.candidateSymbol.fir as? FirSimpleFunction ?: return functionCall
+        lambda.replaceInvocationKind(kind)
-
-        val callsEffects = function.contractDescription.effects
-            ?.map { it.effect }
-            ?.filterIsInstance<ConeCallsEffectDeclaration>() ?: emptyList()
-
-        val isInline = function.isInline
-        if (callsEffects.isEmpty() && !isInline) {
-            return functionCall
-        }
-
-        val reversedArgumentMapping = argumentMapping.entries.map { (argument, parameter) ->
-            parameter to argument
-        }.toMap()
-
-        val invocationKindMapping = mutableMapOf<FirExpression, EventOccurrencesRange>()
-        for (effect in callsEffects) {
-            // TODO: Support callsInPlace contracts on receivers
-            val valueParameter = function.valueParameters.getOrNull(effect.valueParameterReference.parameterIndex) ?: continue
-            val argument = reversedArgumentMapping[valueParameter] ?: continue
-            invocationKindMapping[argument] = effect.kind
-        }
-        if (isInline) {
-            for (argument in functionCall.arguments) {
-                invocationKindMapping.putIfAbsent(argument, EventOccurrencesRange.UNKNOWN)
-            }
-        }
-        if (invocationKindMapping.isEmpty()) {
-            return functionCall
-        }
-        functionCall.argumentList.transformArguments(ArgumentsTransformer, invocationKindMapping to null)
-        return functionCall
     }
 }

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

@@ -31,8 +31,8 @@ import org.jetbrains.kotlin.fir.resolve.dfa.unwrapSmartcastExpression
 import org.jetbrains.kotlin.fir.resolve.diagnostics.*
 import org.jetbrains.kotlin.fir.resolve.inference.FirStubInferenceSession
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
-import org.jetbrains.kotlin.fir.resolve.transformers.InvocationKindTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.StoreReceiver
+import org.jetbrains.kotlin.fir.resolve.transformers.replaceLambdaArgumentInvocationKinds
 import org.jetbrains.kotlin.fir.scopes.impl.isWrappedIntegerOperator
 import org.jetbrains.kotlin.fir.scopes.impl.isWrappedIntegerOperatorForUnsignedType
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
@@ -377,7 +377,7 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         }
         if (calleeReference is FirNamedReferenceWithCandidate) return functionCall
         functionCall.transformAnnotations(transformer, data)
-        functionCall.transformSingle(InvocationKindTransformer, null)
+        functionCall.replaceLambdaArgumentInvocationKinds(session)
         functionCall.transformTypeArguments(transformer, ResolutionMode.ContextIndependent)
         dataFlowAnalyzer.enterFunctionCall(functionCall)
         val (completeInference, callCompleted) =