FIR DFA: update receivers properly in case of single flow in merge point

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -14852,6 +14852,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             runTest("compiler/testData/codegen/box/fir/IrBuiltIns.kt");
         }
 
+        @Test
+        @TestMetadata("LookupTags.kt")
+        public void testLookupTags() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/LookupTags.kt");
+        }
+
         @Test
         @TestMetadata("NameHighlighter.kt")
         public void testNameHighlighter() throws Exception {

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

@@ -37,7 +37,7 @@ class PersistentImplicitReceiverStack private constructor(
 
     fun add(name: Name?, value: ImplicitReceiverValue<*>): PersistentImplicitReceiverStack {
         val stack = stack.add(value)
-        val originalTypes = originalTypes.add(value.type)
+        val originalTypes = originalTypes.add(value.originalType)
         val index = stack.size - 1
         val indexesPerLabel = name?.let { indexesPerLabel.put(it, index) } ?: indexesPerLabel
         val indexesPerSymbol = indexesPerSymbol.put(value.boundSymbol, index)

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

@@ -582,7 +582,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                 syntheticElseNode.mergeIncomingFlow()
             }
         }
-        whenExitNode.mergeIncomingFlow(updateReceivers = true)
+        whenExitNode.mergeIncomingFlow()
     }
 
     // ----------------------------------- While Loop -----------------------------------
@@ -675,7 +675,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     fun enterFinallyBlock() {
         // NB: fork to isolate effects inside the finally block
         // Otherwise, changes in the finally block could affect the previous nodes: try main block and catch clauses.
-        graphBuilder.enterFinallyBlock().mergeIncomingFlow(shouldForkFlow = true)
+        graphBuilder.enterFinallyBlock().mergeIncomingFlow(updateReceivers = true, shouldForkFlow = true)
     }
 
     fun exitFinallyBlock(tryExpression: FirTryExpression) {
@@ -1144,15 +1144,28 @@ 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 ->
-        val previousFlows = if (node.isDead)
-            node.previousNodes.mapNotNull { runIf(!node.incomingEdges.getValue(it).kind.isBack) { it.flow } }
-        else
-            node.previousNodes.mapNotNull { prev -> prev.takeIf { node.incomingEdges.getValue(it).kind.usedInDfa }?.flow }
+        val previousFlows = mutableListOf<FLOW>()
+        var deadForwardCount = 0
+        for (previousNode in previousNodes) {
+            val incomingEdgeKind = node.incomingEdges.getValue(previousNode).kind
+            if (node.isDead) {
+                if (!incomingEdgeKind.isBack) {
+                    previousFlows += previousNode.flow
+                }
+            } else if (incomingEdgeKind.usedInDfa) {
+                previousFlows += previousNode.flow
+            }
+            if (incomingEdgeKind == EdgeKind.DeadForward) {
+                deadForwardCount++
+            }
+        }
         var flow = logicSystem.joinFlow(previousFlows)
-        if (updateReceivers) {
+        // deadForwardCount should be added due to cases like merge after 'if (...) return else ...'
+        if (updateReceivers || previousFlows.size + deadForwardCount > 1) {
             logicSystem.updateAllReceivers(flow)
         }
         if (shouldForkFlow) {

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

@@ -138,9 +138,6 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
 
         commonFlow.addVariableAliases(aliasedVariablesThatDontChangeAlias)
-
-        updateAllReceivers(commonFlow)
-
         return commonFlow
     }
 

compiler/testData/codegen/box/fir/LookupTags.kt

@@ -0,0 +1,39 @@
+// TARGET_BACKEND: JVM
+
+abstract class ConeClassLikeLookupTag {
+    abstract val status: String
+}
+
+class ConeClassLikeLookupTagImpl : ConeClassLikeLookupTag() {
+    override var status: String = ""
+        private set
+
+    fun change(newStatus: String) {
+        status = newStatus
+    }
+}
+
+class ConeClassLikeErrorLookupTag : ConeClassLikeLookupTag() {
+    override val status: String
+        get() = "ERROR"
+}
+
+fun ConeClassLikeLookupTag.foo(): String {
+    (this as? ConeClassLikeLookupTagImpl)?.status?.takeIf { it == "OK" }?.let { return it }
+    return status.also {
+        (this as? ConeClassLikeLookupTagImpl)?.bar(it)
+    }
+}
+
+fun ConeClassLikeLookupTagImpl.bar(s: String) {
+    change(s)
+}
+
+fun box(): String {
+    val tag = ConeClassLikeErrorLookupTag()
+    tag.foo()
+    val tag2 = ConeClassLikeLookupTagImpl()
+    tag2.change("OK")
+    tag2.foo()
+    return tag2.status
+}

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

@@ -14852,6 +14852,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTest("compiler/testData/codegen/box/fir/IrBuiltIns.kt");
         }
 
+        @Test
+        @TestMetadata("LookupTags.kt")
+        public void testLookupTags() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/LookupTags.kt");
+        }
+
         @Test
         @TestMetadata("NameHighlighter.kt")
         public void testNameHighlighter() throws Exception {

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

@@ -14852,6 +14852,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/fir/IrBuiltIns.kt");
         }
 
+        @Test
+        @TestMetadata("LookupTags.kt")
+        public void testLookupTags() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/LookupTags.kt");
+        }
+
         @Test
         @TestMetadata("NameHighlighter.kt")
         public void testNameHighlighter() throws Exception {

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -12283,6 +12283,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/fir/IrBuiltIns.kt");
         }
 
+        @TestMetadata("LookupTags.kt")
+        public void testLookupTags() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/LookupTags.kt");
+        }
+
         @TestMetadata("NameHighlighter.kt")
         public void testNameHighlighter() throws Exception {
             runTest("compiler/testData/codegen/box/fir/NameHighlighter.kt");