FIR: fork flow on function entry, restore receivers on exit

^KT-52680 Fixed

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

@@ -28558,6 +28558,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 runTest("compiler/testData/diagnostics/tests/smartCasts/implicitReceiver.kt");
             }
 
+            @Test
+            @TestMetadata("implicitReceiverAcrossFunctions.kt")
+            public void testImplicitReceiverAcrossFunctions() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/smartCasts/implicitReceiverAcrossFunctions.kt");
+            }
+
             @Test
             @TestMetadata("implicitThisOnRefInLambdaInSmartcast.kt")
             public void testImplicitThisOnRefInLambdaInSmartcast() throws Exception {

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

@@ -28558,6 +28558,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 runTest("compiler/testData/diagnostics/tests/smartCasts/implicitReceiver.kt");
             }
 
+            @Test
+            @TestMetadata("implicitReceiverAcrossFunctions.kt")
+            public void testImplicitReceiverAcrossFunctions() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/smartCasts/implicitReceiverAcrossFunctions.kt");
+            }
+
             @Test
             @TestMetadata("implicitThisOnRefInLambdaInSmartcast.kt")
             public void testImplicitThisOnRefInLambdaInSmartcast() throws Exception {

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

@@ -28558,6 +28558,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 runTest("compiler/testData/diagnostics/tests/smartCasts/implicitReceiver.kt");
             }
 
+            @Test
+            @TestMetadata("implicitReceiverAcrossFunctions.kt")
+            public void testImplicitReceiverAcrossFunctions() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/smartCasts/implicitReceiverAcrossFunctions.kt");
+            }
+
             @Test
             @TestMetadata("implicitThisOnRefInLambdaInSmartcast.kt")
             public void testImplicitThisOnRefInLambdaInSmartcast() throws Exception {

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

@@ -269,6 +269,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         if (graphBuilder.isTopLevel()) {
             context.reset()
         }
+        logicSystem.updateAllReceivers(graph.enterNode.computeIncomingFlow().first)
         return FirControlFlowGraphReferenceImpl(graph, DataFlowInfo(variableStorage, flowOnNodes))
     }
 
@@ -282,7 +283,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         val (postponedLambdaEnterNode, functionEnterNode) = graphBuilder.enterAnonymousFunction(anonymousFunction)
         // TODO: questionable
         postponedLambdaEnterNode?.mergeIncomingFlow()
-        functionEnterNode.mergeIncomingFlow()
+        functionEnterNode.mergeIncomingFlow(shouldForkFlow = true)
         logicSystem.updateAllReceivers(functionEnterNode.flow)
     }
 
@@ -294,6 +295,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         // TODO: questionable
         postponedLambdaExitNode?.mergeIncomingFlow()
         functionExitNode.mergeIncomingFlow()
+        logicSystem.updateAllReceivers(graph.enterNode.computeIncomingFlow().first)
         return FirControlFlowGraphReferenceImpl(graph)
     }
 
@@ -1512,16 +1514,12 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     private val CFGNode<*>.origin: CFGNode<*> get() = if (this is StubNode) firstPreviousNode else this
 
-    private fun <T : CFGNode<*>> T.mergeIncomingFlow(
-        // This flag should be set true if we're changing flow branches from one to another (e.g. in when, try->catch)
-        updateReceivers: Boolean = false,
-        shouldForkFlow: Boolean = false
-    ): T = this.also { node ->
+    private fun <T : CFGNode<*>> T.computeIncomingFlow(): Pair<FLOW, Int> {
         val previousFlows = mutableListOf<FLOW>()
         var deadForwardCount = 0
         for (previousNode in previousNodes) {
-            val incomingEdgeKind = node.incomingEdges.getValue(previousNode).kind
-            if (node.isDead) {
+            val incomingEdgeKind = incomingEdges.getValue(previousNode).kind
+            if (isDead) {
                 if (!incomingEdgeKind.isBack) {
                     previousFlows += previousNode.flow
                 }
@@ -1532,9 +1530,16 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                 deadForwardCount++
             }
         }
-        var flow = logicSystem.joinFlow(previousFlows)
-        // deadForwardCount should be added due to cases like merge after 'if (...) return else ...'
-        if (updateReceivers || previousFlows.size + deadForwardCount > 1) {
+        return logicSystem.joinFlow(previousFlows) to (previousFlows.size + deadForwardCount)
+    }
+
+    private fun <T : CFGNode<*>> T.mergeIncomingFlow(
+        // This flag should be set true if we're changing flow branches from one to another (e.g. in when, try->catch)
+        updateReceivers: Boolean = false,
+        shouldForkFlow: Boolean = false
+    ): T = this.also { node ->
+        var (flow, incomingEdges) = computeIncomingFlow()
+        if (updateReceivers || incomingEdges > 1) {
             logicSystem.updateAllReceivers(flow)
         }
         if (shouldForkFlow) {

compiler/testData/diagnostics/tests/smartCasts/implicitReceiverAcrossFunctions.fir.kt

@@ -0,0 +1,24 @@
+interface I {
+    val prop: Int
+}
+
+open class A {
+    fun f1() {
+        this as I
+        prop
+    }
+
+    fun f2() {
+        <!UNRESOLVED_REFERENCE!>prop<!>
+    }
+}
+
+open class B {
+    fun f() {
+        {
+            this as I
+            prop
+        }
+        <!UNRESOLVED_REFERENCE!>prop<!>
+    }
+}

compiler/testData/diagnostics/tests/smartCasts/implicitReceiverAcrossFunctions.kt

@@ -0,0 +1,24 @@
+interface I {
+    val prop: Int
+}
+
+open class A {
+    fun f1() {
+        this as I
+        <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>prop<!>
+    }
+
+    fun f2() {
+        <!UNRESOLVED_REFERENCE!>prop<!>
+    }
+}
+
+open class B {
+    fun f() {
+        {
+            this as I
+            <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>prop<!>
+        }
+        <!UNRESOLVED_REFERENCE!>prop<!>
+    }
+}

compiler/testData/diagnostics/tests/smartCasts/implicitReceiverAcrossFunctions.txt

@@ -0,0 +1,25 @@
+package
+
+public open class A {
+    public constructor A()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public final fun f1(): kotlin.Unit
+    public final fun f2(): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public open class B {
+    public constructor B()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public final fun f(): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public interface I {
+    public abstract val prop: kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/testData/diagnostics/tests/smartCasts/loops/doWhileLiteral.fir.kt

@@ -1,10 +0,0 @@
-fun x(): Boolean { return true }
-
-public fun foo(p: String?): Int {
-    // Exotic variant with unused literal
-    do { ->
-        p!!.length
-    } while (!x())
-    // Literal is not called so p.length is unsafe
-    return p.length
-}

compiler/testData/diagnostics/tests/smartCasts/loops/doWhileLiteral.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 fun x(): Boolean { return true }
 
 public fun foo(p: String?): Int {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -28648,6 +28648,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/tests/smartCasts/implicitReceiver.kt");
             }
 
+            @Test
+            @TestMetadata("implicitReceiverAcrossFunctions.kt")
+            public void testImplicitReceiverAcrossFunctions() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/smartCasts/implicitReceiverAcrossFunctions.kt");
+            }
+
             @Test
             @TestMetadata("implicitThisOnRefInLambdaInSmartcast.kt")
             public void testImplicitThisOnRefInLambdaInSmartcast() throws Exception {

compiler/tests-spec/testData/diagnostics/notLinked/dfa/pos/53.fir.kt

@@ -19,8 +19,8 @@ fun case_2(x: Int?) = 10
 fun case_2() {
     var x: Int? = 10
     var y = { x = null }
-    if (<!SENSELESS_COMPARISON!>x != null<!>) {
-        val z = case_2(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int? & kotlin.Nothing")!>x<!>)
+    if (x != null) {
+        val z = case_2(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int? & kotlin.Int")!>x<!>)
         <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>z<!>
     }
 }