FIR CFG: merge data flow if called-in-place lambda may not be called

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

@@ -1105,6 +1105,12 @@ public class DiagnosisCompilerFirTestdataTestGenerated extends AbstractDiagnosis
                 runTest("compiler/fir/analysis-tests/testData/resolve/cfg/flowFromInplaceLambda2.kt");
             }
 
+            @Test
+            @TestMetadata("flowFromInplaceLambda3.kt")
+            public void testFlowFromInplaceLambda3() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/cfg/flowFromInplaceLambda3.kt");
+            }
+
             @Test
             @TestMetadata("flowFromTwoInplaceLambdas.kt")
             public void testFlowFromTwoInplaceLambdas() throws Exception {

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

@@ -956,6 +956,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
             runTest("compiler/fir/analysis-tests/testData/resolve/cfg/flowFromInplaceLambda2.kt");
         }
 
+        @TestMetadata("flowFromInplaceLambda3.kt")
+        public void testFlowFromInplaceLambda3() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/cfg/flowFromInplaceLambda3.kt");
+        }
+
         @TestMetadata("flowFromTwoInplaceLambdas.kt")
         public void testFlowFromTwoInplaceLambdas() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/cfg/flowFromTwoInplaceLambdas.kt");

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

@@ -0,0 +1,307 @@
+digraph flowFromInplaceLambda3_kt {
+    graph [nodesep=3]
+    node [shape=box penwidth=2]
+    edge [penwidth=2]
+
+    subgraph cluster_0 {
+        color=red
+        0 [label="Enter function unknown" style="filled" fillcolor=red];
+        subgraph cluster_1 {
+            color=blue
+            1 [label="Enter block"];
+            2 [label="Function call: R|<local>/x|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()"];
+            3 [label="Exit block"];
+        }
+        4 [label="Exit function unknown" style="filled" fillcolor=red];
+    }
+    0 -> {1};
+    1 -> {2};
+    2 -> {3};
+    3 -> {4};
+
+    subgraph cluster_2 {
+        color=red
+        5 [label="Enter function atLeastOnce" style="filled" fillcolor=red];
+        subgraph cluster_3 {
+            color=blue
+            6 [label="Enter block"];
+            7 [label="Function call: R|<local>/x|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()"];
+            8 [label="Exit block"];
+        }
+        9 [label="Exit function atLeastOnce" style="filled" fillcolor=red];
+    }
+    5 -> {6};
+    6 -> {7};
+    7 -> {8};
+    8 -> {9};
+
+    subgraph cluster_4 {
+        color=red
+        10 [label="Enter function exactlyOnce" style="filled" fillcolor=red];
+        subgraph cluster_5 {
+            color=blue
+            11 [label="Enter block"];
+            12 [label="Function call: R|<local>/x|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()"];
+            13 [label="Exit block"];
+        }
+        14 [label="Exit function exactlyOnce" style="filled" fillcolor=red];
+    }
+    10 -> {11};
+    11 -> {12};
+    12 -> {13};
+    13 -> {14};
+
+    subgraph cluster_6 {
+        color=red
+        15 [label="Enter function atMostOnce" style="filled" fillcolor=red];
+        subgraph cluster_7 {
+            color=blue
+            16 [label="Enter block"];
+            17 [label="Exit block"];
+        }
+        18 [label="Exit function atMostOnce" style="filled" fillcolor=red];
+    }
+    15 -> {16};
+    16 -> {17};
+    17 -> {18};
+
+    subgraph cluster_8 {
+        color=red
+        19 [label="Enter function test1" style="filled" fillcolor=red];
+        subgraph cluster_9 {
+            color=blue
+            20 [label="Enter block"];
+            21 [label="Variable declaration: lvar x: R|kotlin/Any?|"];
+            22 [label="Const: String()"];
+            23 [label="Assignment: R|<local>/x|"];
+            24 [label="Access variable R|<local>/x|"];
+            25 [label="Access variable R|kotlin/String.length|"];
+            26 [label="Postponed enter to lambda"];
+            subgraph cluster_10 {
+                color=blue
+                35 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+                subgraph cluster_11 {
+                    color=blue
+                    36 [label="Enter block"];
+                    37 [label="Const: Int(1)"];
+                    38 [label="Assignment: R|<local>/x|"];
+                    39 [label="Exit block"];
+                }
+                40 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+            }
+            27 [label="Postponed exit from lambda"];
+            28 [label="Function call: R|/unknown|(...)"];
+            29 [label="Access variable R|<local>/x|"];
+            30 [label="Access variable <Unresolved name: length>#"];
+            31 [label="Access variable R|<local>/x|"];
+            32 [label="Function call: R|<local>/x|.<Unresolved name: inc>#()"];
+            33 [label="Exit block"];
+        }
+        34 [label="Exit function test1" style="filled" fillcolor=red];
+    }
+    19 -> {20};
+    20 -> {21};
+    21 -> {22};
+    22 -> {23};
+    23 -> {24};
+    24 -> {25};
+    25 -> {26};
+    26 -> {27 35};
+    26 -> {35} [style=dashed];
+    27 -> {28};
+    28 -> {29};
+    29 -> {30};
+    30 -> {31};
+    31 -> {32};
+    32 -> {33};
+    33 -> {34};
+    35 -> {40 36};
+    36 -> {37};
+    37 -> {38};
+    38 -> {39};
+    39 -> {40};
+    40 -> {27};
+    40 -> {35} [color=green style=dashed];
+
+    subgraph cluster_12 {
+        color=red
+        41 [label="Enter function test2" style="filled" fillcolor=red];
+        subgraph cluster_13 {
+            color=blue
+            42 [label="Enter block"];
+            43 [label="Variable declaration: lvar x: R|kotlin/Any?|"];
+            44 [label="Const: String()"];
+            45 [label="Assignment: R|<local>/x|"];
+            46 [label="Access variable R|<local>/x|"];
+            47 [label="Access variable R|kotlin/String.length|"];
+            48 [label="Postponed enter to lambda"];
+            subgraph cluster_14 {
+                color=blue
+                58 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+                subgraph cluster_15 {
+                    color=blue
+                    59 [label="Enter block"];
+                    60 [label="Const: Int(1)"];
+                    61 [label="Assignment: R|<local>/x|"];
+                    62 [label="Exit block"];
+                }
+                63 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+            }
+            49 [label="Call arguments union" style="filled" fillcolor=yellow];
+            50 [label="Postponed exit from lambda"];
+            51 [label="Function call: R|/atLeastOnce|(...)"];
+            52 [label="Access variable R|<local>/x|"];
+            53 [label="Access variable <Unresolved name: length>#"];
+            54 [label="Access variable R|<local>/x|"];
+            55 [label="Function call: R|<local>/x|.R|kotlin/Int.inc|()"];
+            56 [label="Exit block"];
+        }
+        57 [label="Exit function test2" style="filled" fillcolor=red];
+    }
+    41 -> {42};
+    42 -> {43};
+    43 -> {44};
+    44 -> {45};
+    45 -> {46};
+    46 -> {47};
+    47 -> {48};
+    48 -> {58};
+    48 -> {50} [color=red];
+    48 -> {58} [style=dashed];
+    49 -> {51} [color=red];
+    50 -> {51} [color=green];
+    51 -> {52};
+    52 -> {53};
+    53 -> {54};
+    54 -> {55};
+    55 -> {56};
+    56 -> {57};
+    58 -> {59};
+    59 -> {60};
+    60 -> {61};
+    61 -> {62};
+    62 -> {63};
+    63 -> {49} [color=red];
+    63 -> {50} [color=green];
+    63 -> {58} [color=green style=dashed];
+
+    subgraph cluster_16 {
+        color=red
+        64 [label="Enter function test3" style="filled" fillcolor=red];
+        subgraph cluster_17 {
+            color=blue
+            65 [label="Enter block"];
+            66 [label="Variable declaration: lvar x: R|kotlin/Any?|"];
+            67 [label="Const: String()"];
+            68 [label="Assignment: R|<local>/x|"];
+            69 [label="Access variable R|<local>/x|"];
+            70 [label="Access variable R|kotlin/String.length|"];
+            71 [label="Postponed enter to lambda"];
+            subgraph cluster_18 {
+                color=blue
+                81 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+                subgraph cluster_19 {
+                    color=blue
+                    82 [label="Enter block"];
+                    83 [label="Const: Int(1)"];
+                    84 [label="Assignment: R|<local>/x|"];
+                    85 [label="Exit block"];
+                }
+                86 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+            }
+            72 [label="Call arguments union" style="filled" fillcolor=yellow];
+            73 [label="Postponed exit from lambda"];
+            74 [label="Function call: R|/exactlyOnce|(...)"];
+            75 [label="Access variable R|<local>/x|"];
+            76 [label="Access variable <Unresolved name: length>#"];
+            77 [label="Access variable R|<local>/x|"];
+            78 [label="Function call: R|<local>/x|.R|kotlin/Int.inc|()"];
+            79 [label="Exit block"];
+        }
+        80 [label="Exit function test3" style="filled" fillcolor=red];
+    }
+    64 -> {65};
+    65 -> {66};
+    66 -> {67};
+    67 -> {68};
+    68 -> {69};
+    69 -> {70};
+    70 -> {71};
+    71 -> {81};
+    71 -> {73} [color=red];
+    71 -> {81} [style=dashed];
+    72 -> {74} [color=red];
+    73 -> {74} [color=green];
+    74 -> {75};
+    75 -> {76};
+    76 -> {77};
+    77 -> {78};
+    78 -> {79};
+    79 -> {80};
+    81 -> {82};
+    82 -> {83};
+    83 -> {84};
+    84 -> {85};
+    85 -> {86};
+    86 -> {72} [color=red];
+    86 -> {73} [color=green];
+    86 -> {81} [color=green style=dashed];
+
+    subgraph cluster_20 {
+        color=red
+        87 [label="Enter function test4" style="filled" fillcolor=red];
+        subgraph cluster_21 {
+            color=blue
+            88 [label="Enter block"];
+            89 [label="Variable declaration: lvar x: R|kotlin/Any?|"];
+            90 [label="Const: String()"];
+            91 [label="Assignment: R|<local>/x|"];
+            92 [label="Access variable R|<local>/x|"];
+            93 [label="Access variable R|kotlin/String.length|"];
+            94 [label="Postponed enter to lambda"];
+            subgraph cluster_22 {
+                color=blue
+                103 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+                subgraph cluster_23 {
+                    color=blue
+                    104 [label="Enter block"];
+                    105 [label="Const: Int(1)"];
+                    106 [label="Assignment: R|<local>/x|"];
+                    107 [label="Exit block"];
+                }
+                108 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+            }
+            95 [label="Postponed exit from lambda"];
+            96 [label="Function call: R|/atMostOnce|(...)"];
+            97 [label="Access variable R|<local>/x|"];
+            98 [label="Access variable <Unresolved name: length>#"];
+            99 [label="Access variable R|<local>/x|"];
+            100 [label="Function call: R|<local>/x|.<Unresolved name: inc>#()"];
+            101 [label="Exit block"];
+        }
+        102 [label="Exit function test4" style="filled" fillcolor=red];
+    }
+    87 -> {88};
+    88 -> {89};
+    89 -> {90};
+    90 -> {91};
+    91 -> {92};
+    92 -> {93};
+    93 -> {94};
+    94 -> {95 103};
+    94 -> {103} [style=dashed];
+    95 -> {96};
+    96 -> {97};
+    97 -> {98};
+    98 -> {99};
+    99 -> {100};
+    100 -> {101};
+    101 -> {102};
+    103 -> {108 104};
+    104 -> {105};
+    105 -> {106};
+    106 -> {107};
+    107 -> {108};
+    108 -> {95};
+
+}

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

@@ -0,0 +1,80 @@
+FILE: flowFromInplaceLambda3.kt
+    public final fun unknown(x: R|() -> kotlin/Unit|): R|kotlin/Unit|
+        [R|Contract description]
+         <
+            CallsInPlace(x, UNKNOWN)
+        >
+     {
+        [StubStatement]
+        R|<local>/x|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()
+    }
+    public final fun atLeastOnce(x: R|() -> kotlin/Unit|): R|kotlin/Unit|
+        [R|Contract description]
+         <
+            CallsInPlace(x, AT_LEAST_ONCE)
+        >
+     {
+        [StubStatement]
+        R|<local>/x|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()
+    }
+    public final fun exactlyOnce(x: R|() -> kotlin/Unit|): R|kotlin/Unit|
+        [R|Contract description]
+         <
+            CallsInPlace(x, AT_LEAST_ONCE)
+        >
+     {
+        [StubStatement]
+        R|<local>/x|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()
+    }
+    public final fun atMostOnce(x: R|() -> kotlin/Unit|): R|kotlin/Unit|
+        [R|Contract description]
+         <
+            CallsInPlace(x, AT_MOST_ONCE)
+        >
+     {
+        [StubStatement]
+    }
+    public final fun test1(): R|kotlin/Unit| {
+        lvar x: R|kotlin/Any?|
+        R|<local>/x| = String()
+        R|<local>/x|.R|kotlin/String.length|
+        R|/unknown|(<L> = unknown@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline, kind=UNKNOWN>  {
+            R|<local>/x| = Int(1)
+        }
+        )
+        R|<local>/x|.<Unresolved name: length>#
+        R|<local>/x|.<Unresolved name: inc>#()
+    }
+    public final fun test2(): R|kotlin/Unit| {
+        lvar x: R|kotlin/Any?|
+        R|<local>/x| = String()
+        R|<local>/x|.R|kotlin/String.length|
+        R|/atLeastOnce|(<L> = atLeastOnce@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline, kind=AT_LEAST_ONCE>  {
+            R|<local>/x| = Int(1)
+        }
+        )
+        R|<local>/x|.<Unresolved name: length>#
+        R|<local>/x|.R|kotlin/Int.inc|()
+    }
+    public final fun test3(): R|kotlin/Unit| {
+        lvar x: R|kotlin/Any?|
+        R|<local>/x| = String()
+        R|<local>/x|.R|kotlin/String.length|
+        R|/exactlyOnce|(<L> = exactlyOnce@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline, kind=AT_LEAST_ONCE>  {
+            R|<local>/x| = Int(1)
+        }
+        )
+        R|<local>/x|.<Unresolved name: length>#
+        R|<local>/x|.R|kotlin/Int.inc|()
+    }
+    public final fun test4(): R|kotlin/Unit| {
+        lvar x: R|kotlin/Any?|
+        R|<local>/x| = String()
+        R|<local>/x|.R|kotlin/String.length|
+        R|/atMostOnce|(<L> = atMostOnce@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline, kind=AT_MOST_ONCE>  {
+            R|<local>/x| = Int(1)
+        }
+        )
+        R|<local>/x|.<Unresolved name: length>#
+        R|<local>/x|.<Unresolved name: inc>#()
+    }

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

@@ -0,0 +1,58 @@
+// !DUMP_CFG
+
+import kotlin.contracts.*
+
+fun unknown(x: () -> Unit) {
+    contract { callsInPlace(x, InvocationKind.UNKNOWN) }
+    x()
+}
+
+fun atLeastOnce(x: () -> Unit) {
+    contract { callsInPlace(x, InvocationKind.AT_LEAST_ONCE) }
+    x()
+}
+
+fun exactlyOnce(x: () -> Unit) {
+    contract { callsInPlace(x, InvocationKind.AT_LEAST_ONCE) }
+    x()
+}
+
+fun atMostOnce(x: () -> Unit) {
+    contract { callsInPlace(x, InvocationKind.AT_MOST_ONCE) }
+}
+
+fun test1() {
+    var x: Any?
+    x = ""
+    x.length
+    unknown { x = 1 }
+    x.<!UNRESOLVED_REFERENCE!>length<!>
+    x.<!UNRESOLVED_REFERENCE!>inc<!>()
+}
+
+fun test2() {
+    var x: Any?
+    x = ""
+    x.length
+    atLeastOnce { x = 1 }
+    x.<!UNRESOLVED_REFERENCE!>length<!>
+    x.inc()
+}
+
+fun test3() {
+    var x: Any?
+    x = ""
+    x.length
+    exactlyOnce { x = 1 }
+    x.<!UNRESOLVED_REFERENCE!>length<!>
+    x.inc()
+}
+
+fun test4() {
+    var x: Any?
+    x = ""
+    x.length
+    atMostOnce { x = 1 }
+    x.<!UNRESOLVED_REFERENCE!>length<!>
+    x.<!UNRESOLVED_REFERENCE!>inc<!>()
+}

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

@@ -395,8 +395,7 @@ digraph jumps_kt {
     }
     124 -> {125};
     125 -> {126};
-    126 -> {131};
+    126 -> {127 131};
-    126 -> {127} [color=red];
     126 -> {131} [style=dashed];
     127 -> {128};
     128 -> {129};
@@ -407,7 +406,7 @@ digraph jumps_kt {
     133 -> {134} [style=dotted];
     134 -> {135} [style=dotted];
     135 -> {136} [style=dotted];
-    136 -> {127} [color=green];
+    136 -> {127};
     136 -> {131} [color=green style=dashed];
 
 }

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

@@ -74,8 +74,7 @@ digraph lambdas_kt {
     12 -> {20};
     13 -> {14};
     14 -> {15};
-    15 -> {23};
+    15 -> {16 23};
-    15 -> {16} [color=red];
     15 -> {23} [style=dashed];
     16 -> {17};
     17 -> {18};
@@ -88,7 +87,7 @@ digraph lambdas_kt {
     25 -> {26};
     26 -> {27};
     27 -> {28};
-    28 -> {16} [color=green];
+    28 -> {16};
     28 -> {23} [color=green style=dashed];
 
     subgraph cluster_9 {
@@ -218,8 +217,7 @@ digraph lambdas_kt {
     }
     62 -> {63};
     63 -> {64};
-    64 -> {71};
+    64 -> {65 71};
-    64 -> {65} [color=red];
     64 -> {71} [style=dashed];
     65 -> {66};
     66 -> {67};
@@ -234,7 +232,7 @@ digraph lambdas_kt {
     74 -> {75} [style=dotted];
     75 -> {76} [style=dotted];
     76 -> {77} [style=dotted];
-    77 -> {65} [color=green];
+    77 -> {65};
     77 -> {71} [color=green style=dashed];
 
     subgraph cluster_22 {
@@ -270,8 +268,7 @@ digraph lambdas_kt {
     }
     78 -> {79};
     79 -> {80};
-    80 -> {87};
+    80 -> {81 87};
-    80 -> {81} [color=red];
     80 -> {87} [style=dashed];
     81 -> {82};
     82 -> {83};
@@ -286,7 +283,7 @@ digraph lambdas_kt {
     90 -> {91} [style=dotted];
     91 -> {92} [style=dotted];
     92 -> {93} [style=dotted];
-    93 -> {81} [color=green];
+    93 -> {81};
     93 -> {87} [color=green style=dashed];
 
 }

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

@@ -191,8 +191,7 @@ digraph localClassesWithImplicit_kt {
     }
     36 -> {37};
     37 -> {38};
-    38 -> {45};
+    38 -> {39 45};
-    38 -> {39} [color=red];
     38 -> {45} [style=dashed];
     39 -> {40};
     40 -> {41};
@@ -225,7 +224,7 @@ digraph localClassesWithImplicit_kt {
     67 -> {68};
     68 -> {69};
     69 -> {70};
-    70 -> {39} [color=green];
+    70 -> {39};
     70 -> {45} [color=green style=dashed];
 
     subgraph cluster_19 {
@@ -265,8 +264,7 @@ digraph localClassesWithImplicit_kt {
     }
     71 -> {72};
     72 -> {73};
-    73 -> {80};
+    73 -> {74 80};
-    73 -> {74} [color=red];
     73 -> {80} [style=dashed];
     74 -> {75};
     75 -> {76};
@@ -282,7 +280,7 @@ digraph localClassesWithImplicit_kt {
     85 -> {86};
     86 -> {87};
     87 -> {88};
-    88 -> {74} [color=green];
+    88 -> {74};
     88 -> {80} [color=green style=dashed];
 
     subgraph cluster_23 {
@@ -392,8 +390,7 @@ digraph localClassesWithImplicit_kt {
     }
     99 -> {100};
     100 -> {101};
-    101 -> {108};
+    101 -> {102 108};
-    101 -> {102} [color=red];
     101 -> {108} [style=dashed];
     102 -> {103};
     103 -> {104};
@@ -426,7 +423,7 @@ digraph localClassesWithImplicit_kt {
     130 -> {131};
     131 -> {132};
     132 -> {133};
-    133 -> {102} [color=green];
+    133 -> {102};
     133 -> {108} [color=green style=dashed];
 
     subgraph cluster_35 {
@@ -466,8 +463,7 @@ digraph localClassesWithImplicit_kt {
     }
     134 -> {135};
     135 -> {136};
-    136 -> {143};
+    136 -> {137 143};
-    136 -> {137} [color=red];
     136 -> {143} [style=dashed];
     137 -> {138};
     138 -> {139};
@@ -483,7 +479,7 @@ digraph localClassesWithImplicit_kt {
     148 -> {149};
     149 -> {150};
     150 -> {151};
-    151 -> {137} [color=green];
+    151 -> {137};
     151 -> {143} [color=green style=dashed];
 
     subgraph cluster_39 {

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

@@ -47,8 +47,7 @@ digraph postponedLambdas_kt {
     5 -> {6};
     6 -> {7};
     7 -> {8};
-    8 -> {13};
+    8 -> {9 13};
-    8 -> {9} [color=red];
     8 -> {13} [style=dashed];
     9 -> {10};
     10 -> {11};
@@ -57,7 +56,7 @@ digraph postponedLambdas_kt {
     14 -> {15};
     15 -> {16};
     16 -> {17};
-    17 -> {9} [color=green];
+    17 -> {9};
     17 -> {13} [color=green style=dashed];
 
 }

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

@@ -227,8 +227,7 @@ digraph propertiesAndInitBlocks_kt {
         79 [label="Exit property" style="filled" fillcolor=red];
     }
     75 -> {76};
-    76 -> {24};
+    76 -> {77 24};
-    76 -> {77} [color=red];
     76 -> {24} [style=dashed];
     77 -> {78};
     78 -> {79};
@@ -243,7 +242,7 @@ digraph propertiesAndInitBlocks_kt {
     28 -> {29} [style=dotted];
     29 -> {30} [style=dotted];
     30 -> {31} [style=dotted];
-    31 -> {77} [color=green];
+    31 -> {77};
     31 -> {24} [color=green style=dashed];
     32 -> {33} [color=green];
     33 -> {34} [style=dotted];

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

@@ -80,8 +80,7 @@ digraph callsInPlace_kt {
     17 -> {18};
     18 -> {19};
     19 -> {20};
-    20 -> {25};
+    20 -> {21 25};
-    20 -> {21} [color=red];
     20 -> {25} [style=dashed];
     21 -> {22};
     22 -> {23};
@@ -90,7 +89,7 @@ digraph callsInPlace_kt {
     26 -> {27};
     27 -> {28};
     28 -> {29};
-    29 -> {21} [color=green];
+    29 -> {21};
     29 -> {25} [color=green style=dashed];
 
     subgraph cluster_8 {
@@ -120,8 +119,7 @@ digraph callsInPlace_kt {
     }
     30 -> {31};
     31 -> {32};
-    32 -> {38};
+    32 -> {33 38};
-    32 -> {33} [color=red];
     32 -> {38} [style=dashed];
     33 -> {34};
     34 -> {35};
@@ -131,7 +129,7 @@ digraph callsInPlace_kt {
     39 -> {40};
     40 -> {41};
     41 -> {42};
-    42 -> {33} [color=green];
+    42 -> {33};
     42 -> {38} [color=green style=dashed];
 
     subgraph cluster_12 {
@@ -293,12 +291,10 @@ digraph callsInPlace_kt {
     }
     85 -> {86};
     86 -> {87};
-    87 -> {99};
+    87 -> {88 99};
-    87 -> {88} [color=red];
     87 -> {99} [style=dashed];
     88 -> {89};
-    89 -> {94};
+    89 -> {90 94};
-    89 -> {90} [color=red];
     89 -> {94} [style=dashed];
     90 -> {91};
     91 -> {92};
@@ -307,13 +303,13 @@ digraph callsInPlace_kt {
     95 -> {96};
     96 -> {97};
     97 -> {98};
-    98 -> {90} [color=green];
+    98 -> {90};
     98 -> {94} [color=green style=dashed];
     99 -> {103 100};
     100 -> {101};
     101 -> {102};
     102 -> {103};
-    103 -> {88} [color=green];
+    103 -> {88};
     103 -> {99} [color=green style=dashed];
 
     subgraph cluster_28 {
@@ -355,12 +351,10 @@ digraph callsInPlace_kt {
     }
     104 -> {105};
     105 -> {106};
-    106 -> {113};
+    106 -> {107 113};
-    106 -> {107} [color=red];
     106 -> {113} [style=dashed];
     107 -> {108};
-    108 -> {118};
+    108 -> {109 118};
-    108 -> {109} [color=red];
     108 -> {118} [style=dashed];
     109 -> {110};
     110 -> {111};
@@ -369,13 +363,13 @@ digraph callsInPlace_kt {
     114 -> {115};
     115 -> {116};
     116 -> {117};
-    117 -> {107} [color=green];
+    117 -> {107};
     117 -> {113} [color=green style=dashed];
     118 -> {122 119};
     119 -> {120};
     120 -> {121};
     121 -> {122};
-    122 -> {109} [color=green];
+    122 -> {109};
     122 -> {118} [color=green style=dashed];
 
     subgraph cluster_34 {

compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/callsInPlace/atLeastOnce.dot

@@ -45,88 +45,94 @@ digraph atLeastOnce_kt {
             13 [label="Postponed enter to lambda"];
             subgraph cluster_6 {
                 color=blue
-                20 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+                21 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
                 subgraph cluster_7 {
                     color=blue
-                    21 [label="Enter block"];
+                    22 [label="Enter block"];
-                    22 [label="Const: Int(1)"];
+                    23 [label="Const: Int(1)"];
-                    23 [label="Assignment: R|<local>/x|"];
+                    24 [label="Assignment: R|<local>/x|"];
-                    24 [label="Exit block"];
+                    25 [label="Exit block"];
                 }
-                25 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+                26 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
             }
-            14 [label="Postponed exit from lambda"];
+            14 [label="Call arguments union" style="filled" fillcolor=yellow];
-            15 [label="Function call: R|/inlineRun|(...)"];
+            15 [label="Postponed exit from lambda"];
-            16 [label="Access variable R|<local>/x|"];
+            16 [label="Function call: R|/inlineRun|(...)"];
-            17 [label="Function call: R|<local>/x|.R|kotlin/Int.inc|()"];
+            17 [label="Access variable R|<local>/x|"];
-            18 [label="Exit block"];
+            18 [label="Function call: R|<local>/x|.R|kotlin/Int.inc|()"];
+            19 [label="Exit block"];
         }
-        19 [label="Exit function test_1" style="filled" fillcolor=red];
+        20 [label="Exit function test_1" style="filled" fillcolor=red];
     }
     10 -> {11};
     11 -> {12};
     12 -> {13};
-    13 -> {20};
+    13 -> {21};
-    13 -> {14} [color=red];
+    13 -> {15} [color=red];
-    13 -> {20} [style=dashed];
+    13 -> {21} [style=dashed];
-    14 -> {15};
+    14 -> {16} [color=red];
-    15 -> {16};
+    15 -> {16} [color=green];
     16 -> {17};
     17 -> {18};
     18 -> {19};
-    20 -> {21};
+    19 -> {20};
     21 -> {22};
     22 -> {23};
     23 -> {24};
     24 -> {25};
-    25 -> {14} [color=green];
+    25 -> {26};
-    25 -> {20} [color=green style=dashed];
+    26 -> {14} [color=red];
+    26 -> {15} [color=green];
+    26 -> {21} [color=green style=dashed];
 
     subgraph cluster_8 {
         color=red
-        26 [label="Enter function test_2" style="filled" fillcolor=red];
+        27 [label="Enter function test_2" style="filled" fillcolor=red];
         subgraph cluster_9 {
             color=blue
-            27 [label="Enter block"];
+            28 [label="Enter block"];
-            28 [label="Variable declaration: lval x: R|kotlin/Int|"];
+            29 [label="Variable declaration: lval x: R|kotlin/Int|"];
-            29 [label="Postponed enter to lambda"];
+            30 [label="Postponed enter to lambda"];
             subgraph cluster_10 {
                 color=blue
-                36 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+                38 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
                 subgraph cluster_11 {
                     color=blue
-                    37 [label="Enter block"];
+                    39 [label="Enter block"];
-                    38 [label="Const: Int(1)"];
+                    40 [label="Const: Int(1)"];
-                    39 [label="Assignment: R|<local>/x|"];
+                    41 [label="Assignment: R|<local>/x|"];
-                    40 [label="Exit block"];
+                    42 [label="Exit block"];
                 }
-                41 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+                43 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
             }
-            30 [label="Postponed exit from lambda"];
+            31 [label="Call arguments union" style="filled" fillcolor=yellow];
-            31 [label="Function call: R|/myRun|(...)"];
+            32 [label="Postponed exit from lambda"];
-            32 [label="Access variable R|<local>/x|"];
+            33 [label="Function call: R|/myRun|(...)"];
-            33 [label="Function call: R|<local>/x|.R|kotlin/Int.inc|()"];
+            34 [label="Access variable R|<local>/x|"];
-            34 [label="Exit block"];
+            35 [label="Function call: R|<local>/x|.R|kotlin/Int.inc|()"];
+            36 [label="Exit block"];
         }
-        35 [label="Exit function test_2" style="filled" fillcolor=red];
+        37 [label="Exit function test_2" style="filled" fillcolor=red];
     }
-    26 -> {27};
     27 -> {28};
     28 -> {29};
-    29 -> {36};
+    29 -> {30};
-    29 -> {30} [color=red];
+    30 -> {38};
-    29 -> {36} [style=dashed];
+    30 -> {32} [color=red];
-    30 -> {31};
+    30 -> {38} [style=dashed];
-    31 -> {32};
+    31 -> {33} [color=red];
-    32 -> {33};
+    32 -> {33} [color=green];
     33 -> {34};
     34 -> {35};
+    35 -> {36};
     36 -> {37};
-    37 -> {38};
     38 -> {39};
     39 -> {40};
     40 -> {41};
-    41 -> {30} [color=green];
+    41 -> {42};
-    41 -> {36} [color=green style=dashed];
+    42 -> {43};
+    43 -> {31} [color=red];
+    43 -> {32} [color=green];
+    43 -> {38} [color=green style=dashed];
 
 }

compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/callsInPlace/atMostOnce.dot

@@ -66,8 +66,7 @@ digraph atMostOnce_kt {
     10 -> {11};
     11 -> {12};
     12 -> {13};
-    13 -> {20};
+    13 -> {14 20};
-    13 -> {14} [color=red];
     13 -> {20} [style=dashed];
     14 -> {15};
     15 -> {16};
@@ -79,7 +78,7 @@ digraph atMostOnce_kt {
     22 -> {23};
     23 -> {24};
     24 -> {25};
-    25 -> {14} [color=green];
+    25 -> {14};
 
     subgraph cluster_8 {
         color=red
@@ -112,8 +111,7 @@ digraph atMostOnce_kt {
     26 -> {27};
     27 -> {28};
     28 -> {29};
-    29 -> {36};
+    29 -> {30 36};
-    29 -> {30} [color=red];
     29 -> {36} [style=dashed];
     30 -> {31};
     31 -> {32};
@@ -125,6 +123,6 @@ digraph atMostOnce_kt {
     38 -> {39};
     39 -> {40};
     40 -> {41};
-    41 -> {30} [color=green];
+    41 -> {30};
 
 }

compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/callsInPlace/inPlaceLambda.dot

@@ -76,8 +76,7 @@ digraph inPlaceLambda_kt {
     15 -> {16};
     16 -> {17};
     17 -> {18};
-    18 -> {23};
+    18 -> {19 23};
-    18 -> {19} [color=red];
     18 -> {23} [style=dashed];
     19 -> {20};
     20 -> {21};
@@ -86,6 +85,6 @@ digraph inPlaceLambda_kt {
     24 -> {25};
     25 -> {26};
     26 -> {27};
-    27 -> {19} [color=green];
+    27 -> {19};
 
 }

compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/callsInPlace/unknown.dot

@@ -66,8 +66,7 @@ digraph unknown_kt {
     10 -> {11};
     11 -> {12};
     12 -> {13};
-    13 -> {20};
+    13 -> {14 20};
-    13 -> {14} [color=red];
     13 -> {20} [style=dashed];
     14 -> {15};
     15 -> {16};
@@ -79,7 +78,7 @@ digraph unknown_kt {
     22 -> {23};
     23 -> {24};
     24 -> {25};
-    25 -> {14} [color=green];
+    25 -> {14};
     25 -> {20} [color=green style=dashed];
 
     subgraph cluster_8 {
@@ -113,8 +112,7 @@ digraph unknown_kt {
     26 -> {27};
     27 -> {28};
     28 -> {29};
-    29 -> {36};
+    29 -> {30 36};
-    29 -> {30} [color=red];
     29 -> {36} [style=dashed];
     30 -> {31};
     31 -> {32};
@@ -126,7 +124,7 @@ digraph unknown_kt {
     38 -> {39};
     39 -> {40};
     40 -> {41};
-    41 -> {30} [color=green];
+    41 -> {30};
     41 -> {36} [color=green style=dashed];
 
 }

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

@@ -146,8 +146,7 @@ digraph plusAssignWithLambdaInRhs_kt {
     37 -> {38};
     38 -> {39};
     39 -> {40};
-    40 -> {45};
+    40 -> {41 45};
-    40 -> {41} [color=red];
     40 -> {45} [style=dashed];
     41 -> {42};
     42 -> {43};
@@ -157,7 +156,7 @@ digraph plusAssignWithLambdaInRhs_kt {
     47 -> {48};
     48 -> {49};
     49 -> {50};
-    50 -> {41} [color=green];
+    50 -> {41};
     50 -> {45} [color=green style=dashed];
 
 }

compiler/fir/analysis-tests/testData/resolveWithStdlib/smartcasts/tryWithLambdaInside.dot

@@ -105,8 +105,7 @@ finally {
     14 -> {15};
     15 -> {16};
     16 -> {17};
-    17 -> {31};
+    17 -> {18 31};
-    17 -> {18} [color=red];
     17 -> {31} [style=dashed];
     18 -> {19};
     19 -> {20};
@@ -126,7 +125,7 @@ finally {
     32 -> {33};
     33 -> {34};
     34 -> {35};
-    35 -> {18} [color=green];
+    35 -> {18};
     35 -> {31} [color=green style=dashed];
 
     subgraph cluster_13 {

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

@@ -1105,6 +1105,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
                 runTest("compiler/fir/analysis-tests/testData/resolve/cfg/flowFromInplaceLambda2.kt");
             }
 
+            @Test
+            @TestMetadata("flowFromInplaceLambda3.kt")
+            public void testFlowFromInplaceLambda3() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/cfg/flowFromInplaceLambda3.kt");
+            }
+
             @Test
             @TestMetadata("flowFromTwoInplaceLambdas.kt")
             public void testFlowFromTwoInplaceLambdas() throws Exception {

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

@@ -1105,6 +1105,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
                 runTest("compiler/fir/analysis-tests/testData/resolve/cfg/flowFromInplaceLambda2.kt");
             }
 
+            @Test
+            @TestMetadata("flowFromInplaceLambda3.kt")
+            public void testFlowFromInplaceLambda3() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/cfg/flowFromInplaceLambda3.kt");
+            }
+
             @Test
             @TestMetadata("flowFromTwoInplaceLambdas.kt")
             public void testFlowFromTwoInplaceLambdas() throws Exception {

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

@@ -308,8 +308,8 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             else -> {}
         }
         // TODO: questionable
-        postponedLambdaExitNode?.mergeIncomingFlow()
         functionExitNode.mergeIncomingFlow()
+        postponedLambdaExitNode?.mergeIncomingFlow()
         logicSystem.updateAllReceivers(graph.enterNode.computeIncomingFlow().first)
         return FirControlFlowGraphReferenceImpl(graph)
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -323,7 +323,7 @@ class ControlFlowGraphBuilder {
 
     private val EventOccurrencesRange?.hasBackEdge: Boolean
         get() = when (this) {
-            EventOccurrencesRange.AT_LEAST_ONCE, EventOccurrencesRange.UNKNOWN -> true
+            EventOccurrencesRange.AT_LEAST_ONCE, EventOccurrencesRange.MORE_THAN_ONCE, EventOccurrencesRange.UNKNOWN -> true
             else -> false
         }
 
@@ -354,20 +354,35 @@ class ControlFlowGraphBuilder {
         }
 
         val invocationKind = anonymousFunction.invocationKind
-        if (invocationKind != null) {
+        val maybeNonZero = invocationKind != null && invocationKind != EventOccurrencesRange.ZERO
-            addEdge(exitNode, postponedExitNode, preferredKind = EdgeKind.CfgForward)
+        val maybeZero = invocationKind != EventOccurrencesRange.AT_LEAST_ONCE &&
-        } else {
+                invocationKind != EventOccurrencesRange.EXACTLY_ONCE &&
+                invocationKind != EventOccurrencesRange.MORE_THAN_ONCE
+        // Four cases we handle differently here:
+        //  1. function not called in-place: data and control flow skips the function,
+        //     and any assignment inside the function invalidates smart casts
+        //  2. function never executed (EventOccurrencesRange.ZERO): same as above,
+        //     but without smart cast invalidation
+        //  3. function executed at least once (or exactly once or more than once):
+        //     control flow merged from postponedEnterNode and exitNode into postponedExitNode,
+        //     data flow goes from exitNode to the union of arguments of the call
+        //  4. function executed an unknown number of times (maybe zero, maybe not):
+        //     both control flow and data flow merged from postponedEnterNode and exitNode
+        //     into postponedExitNode
+        if (maybeZero) {
             val kind = if (postponedExitNode.isDead) EdgeKind.DeadForward else EdgeKind.CfgForward
             CFGNode.addJustKindEdge(postponedEnterNode, postponedExitNode, kind, propagateDeadness = true)
         }
-
+        if (maybeNonZero) {
-        if (invocationKind == EventOccurrencesRange.EXACTLY_ONCE && shouldPassFlowFromInplaceLambda.top()) {
+            addEdge(exitNode, postponedExitNode, preferredKind = if (maybeZero) EdgeKind.Forward else EdgeKind.CfgForward)
-            exitsFromCompletedPostponedAnonymousFunctions.lastOrNull()?.add(postponedExitNode)
+            if (!maybeZero && shouldPassFlowFromInplaceLambda.top()) {
+                exitsFromCompletedPostponedAnonymousFunctions.lastOrNull()?.add(postponedExitNode)
+            }
         }
 
         val containingGraph = parentGraphForAnonymousFunctions.remove(symbol) ?: currentGraph
         containingGraph.addSubGraph(graph)
-        return if (lambdaIsPostponedFromCall) {
+        return if (lambdaIsPostponedFromCall && !(maybeNonZero && maybeZero)) {
             Triple(exitNode, null, graph)
         } else {
             Triple(exitNode, postponedExitNode, graph)

compiler/testData/codegen/box/smartCasts/kt44814.dot

@@ -713,8 +713,7 @@ digraph kt44814_kt {
     188 -> {189};
     189 -> {190};
     190 -> {191};
-    191 -> {239};
+    191 -> {192 239};
-    191 -> {192} [color=red];
     191 -> {239} [style=dashed];
     192 -> {193};
     193 -> {194};
@@ -776,7 +775,7 @@ digraph kt44814_kt {
     246 -> {247};
     247 -> {248};
     248 -> {249};
-    249 -> {192} [color=green];
+    249 -> {192};
     249 -> {239} [color=green style=dashed];
 
     subgraph cluster_68 {

compiler/tests-spec/testData/diagnostics/notLinked/contracts/analysis/smartcasts/pos/12.fir.kt

@@ -10,7 +10,7 @@ fun case_1(arg: Int?) {
 // TESTCASE NUMBER: 2
 fun case_2(arg: Int?) {
     funWithAtLeastOnceCallsInPlace { arg!! }
-    arg<!UNSAFE_CALL!>.<!>inc()
+    arg.inc()
 }
 
 // TESTCASE NUMBER: 3
@@ -24,5 +24,5 @@ fun case_3() {
 fun case_4() {
     val value_1: Boolean?
     funWithAtLeastOnceCallsInPlace { <!VAL_REASSIGNMENT!>value_1<!> = true }
-    value_1<!UNSAFE_CALL!>.<!>not()
+    value_1.not()
 }