[FIR] Make sure not to leak alternate flows outside finally blocks

Instead of trying to use node and edge information to determine when an
alternate flow no longer needs to be propagated, ask the CFG builder if
the node is still within the `finally` blocks of the paths being
propagated. This makes the checks simpler and more sound, avoiding leaks
of alternate flows beyond their needed scope.

KT-56888

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

@@ -644,8 +644,8 @@ abstract class FirDataFlowAnalyzer(
         // there at the time its flow was computed - which is why we erased all information about `possiblyChangedVariables`
         // from it. Now that we have those edges, we can restore type information for the code after the loop.
         val conditionEnterFlow = conditionEnterNode.getFlow(path)
-        val loopEnterAndContinueFlows = conditionEnterNode.previousLiveNodes.map { it.getFlow(path) }.toList()
+        val loopEnterAndContinueFlows = conditionEnterNode.previousLiveNodes.map { it.getFlow(path) }
-        val conditionExitAndBreakFlows = node.previousLiveNodes.map { it.getFlow(path) }.toList()
+        val conditionExitAndBreakFlows = node.previousLiveNodes.map { it.getFlow(path) }
         possiblyChangedVariables.forEach { variable ->
             // The statement about `variable` in `conditionEnterFlow` should be empty, so to obtain the new statement
             // we can simply add the now-known input to whatever was inferred from nothing so long as the value is the same.
@@ -737,7 +737,9 @@ abstract class FirDataFlowAnalyzer(
     }
 
     fun enterFinallyBlock() {
-        graphBuilder.enterFinallyBlock().mergeIncomingFlow()
+        val node = graphBuilder.enterFinallyBlock()
+        node.mergeIncomingFlow()
+        node.createAlternateFlows()
     }
 
     fun exitFinallyBlock() {
@@ -1170,45 +1172,27 @@ abstract class FirDataFlowAnalyzer(
     // receiver stack also correspond to the data flow information attached to `graphBuilder.lastNode`.
     private var currentReceiverState: Flow? = null
 
-    // Generally when calling some method on `graphBuilder`, one of the nodes it returns is the new `lastNode`.
+    private fun CFGNode<*>.buildDefaultFlow(
-    // In that case `mergeIncomingFlow` will automatically ensure consistency once called on that node.
-    private fun CFGNode<*>.buildIncomingFlow(
-        path: FlowPath,
         builder: (FlowPath, MutableFlow) -> Unit,
     ): MutableFlow {
-        val previousFlows = previousDfaNodes.mapNotNull { (edge, node) ->
+        val previousFlows = previousNodes.mapNotNull { node ->
-            // For CFGNodes that cause alternate flow paths to be created, only edges with matching labels should be merged. However, when
+            val edge = edgeFrom(node)
-            // an alternate flow is being propagated through one of these CFGNodes - i.e., when the FirElements do not match - only
+            if (!usedInDfa(edge)) return@mapNotNull null
-            // NormalPath edges should be merged.
-            if (this is AlternateFlowStartMarker && path is FlowPath.CfgEdge) {
-                if (path.fir == this.fir && edge.label != path.label) {
-                    return@mapNotNull null
-                } else if (path.fir != this.fir && edge.label != NormalPath) {
-                    return@mapNotNull null
-                }
-            }
 
             // `MergePostponedLambdaExitsNode` nodes form a parallel data flow graph. We never compute
             // data flow for any of them until reaching a completed call.
             if (node is MergePostponedLambdaExitsNode && !node.flowInitialized) node.mergeIncomingFlow()
 
-            when (path) {
+            // For CFGNodes that are the end of alternate flows, use the alternate flow associated with the edge label.
-                FlowPath.Default -> {
+            if (node is FinallyBlockExitNode) {
-                    // For CFGNodes that are the end of alternate flows, use the alternate flow associated with the edge label.
+                val alternatePath = FlowPath.CfgEdge(edge.label, node.fir)
-                    if (node is AlternateFlowEndMarker) {
+                node.getAlternateFlow(alternatePath) ?: node.flow
-                        val alternatePath = FlowPath.CfgEdge(edge.label, node.fir)
+            } else {
-                        node.getAlternateFlow(alternatePath) ?: node.flow
+                node.flow
-                    } else {
-                        node.flow
-                    }
-                }
-                else -> {
-                    node.getAlternateFlow(path) ?: node.flow
-                }
             }
-        }.toList()
+        }
         val result = logicSystem.joinFlow(previousFlows, isUnion)
-        if (path == FlowPath.Default && graphBuilder.lastNodeOrNull == this) {
+        if (graphBuilder.lastNodeOrNull == this) {
             // Here it is, the new `lastNode`. If the previous state is the only predecessor, then there is actually
             // nothing to update; `addTypeStatement` has already ensured we have the correct information.
             if (currentReceiverState == null || previousFlows.singleOrNull() != currentReceiverState) {
@@ -1216,40 +1200,87 @@ abstract class FirDataFlowAnalyzer(
             }
             currentReceiverState = result
         }
+        return result.also { builder(FlowPath.Default, it) }
+    }
+
+    private fun CFGNode<*>.buildAlternateFlow(
+        path: FlowPath.CfgEdge,
+        builder: (FlowPath, MutableFlow) -> Unit,
+    ): MutableFlow {
+        val alternateFlowStart = this is FinallyBlockEnterNode
+        val previousFlows = previousNodes.mapNotNull { node ->
+            val edge = edgeFrom(node)
+            if (!usedInDfa(edge)) return@mapNotNull null
+
+            // For CFGNodes that cause alternate flow paths to be created, only edges with matching labels should be merged. However, when
+            // an alternate flow is being propagated through one of these CFGNodes - i.e., when the FirElements do not match - only
+            // NormalPath edges should be merged.
+            if (alternateFlowStart) {
+                if (path.fir == this.fir && edge.label != path.label) {
+                    return@mapNotNull null
+                } else if (path.fir != this.fir && edge.label != NormalPath) {
+                    return@mapNotNull null
+                }
+            }
+
+            node.getAlternateFlow(path) ?: node.flow
+        }
+        val result = logicSystem.joinFlow(previousFlows, isUnion)
         return result.also { builder(path, it) }
     }
 
+    // Generally when calling some method on `graphBuilder`, one of the nodes it returns is the new `lastNode`.
+    // In that case `mergeIncomingFlow` will automatically ensure consistency once called on that node.
     @OptIn(CfgInternals::class)
     private fun CFGNode<*>.mergeIncomingFlow(
         builder: (FlowPath, MutableFlow) -> Unit = { _, _ -> },
     ) {
         // Always build the default flow path for all nodes.
-        val mutableDefaultFlow = buildIncomingFlow(FlowPath.Default, builder)
+        val mutableDefaultFlow = buildDefaultFlow(builder)
         val defaultFlow = mutableDefaultFlow.freeze().also { this.flow = it }
         if (currentReceiverState === mutableDefaultFlow) {
             currentReceiverState = defaultFlow
         }
 
         // Propagate alternate flows from previous nodes.
-        val propagatePaths = previousDfaNodes.flatMapTo(mutableSetOf()) { (edge, node) ->
+        propagateAlternateFlows(builder)
-            when (node) {
+    }
+
+    @OptIn(CfgInternals::class)
+    private fun CFGNode<*>.propagateAlternateFlows(
+        builder: (FlowPath, MutableFlow) -> Unit,
+    ) {
+        val propagatedPaths = mutableSetOf<FlowPath>()
+        for (node in previousNodes) {
+            if (node.alternateFlowPaths.isEmpty()) continue
+
+            val edge = edgeFrom(node)
+            // Only propagate alternate flows which originate along a normal path edge and are used in DFA.
+            if (edge.label != NormalPath || !usedInDfa(edge)) continue
+
+            for (path in node.alternateFlowPaths) {
                 // If the source node is the end of alternate flows, do not propagate the alternate flows which have ended.
-                is AlternateFlowEndMarker -> node.alternateFlowPaths.filter { it !is FlowPath.CfgEdge || it.fir != node.fir }
+                if (path !is FlowPath.CfgEdge || !graphBuilder.withinFinallyBlock(path.fir)) continue
-                // Otherwise, only propagate alternate flows which originate along a normal path edge.
+
-                else -> node.alternateFlowPaths.takeIf { edge.label == NormalPath } ?: emptyList()
+                if (propagatedPaths.add(path)) {
+                    addAlternateFlow(path, buildAlternateFlow(path, builder).freeze())
+                }
             }
         }
-        for (path in propagatePaths) {
+    }
-            addAlternateFlow(path, buildIncomingFlow(path, builder).freeze())
-        }
 
-        // Add any new alternate flows that should be created.
+    @OptIn(CfgInternals::class)
-        if (this is AlternateFlowStartMarker) {
+    private fun CFGNode<*>.createAlternateFlows(
-            val additionalPaths = previousDfaNodes
+        builder: (FlowPath, MutableFlow) -> Unit = { _, _ -> },
-                .mapNotNullTo(mutableSetOf()) { (edge, _) -> edge.label.takeIf { it != UncaughtExceptionPath } }
+    ) {
-                .map { FlowPath.CfgEdge(it, this.fir) }
+        val createdLabels = mutableSetOf<EdgeLabel>()
-            for (path in additionalPaths) {
+        for (node in previousNodes) {
-                addAlternateFlow(path, buildIncomingFlow(path, builder).freeze())
+            val edge = edgeFrom(node)
+            if (edge.label == UncaughtExceptionPath || !usedInDfa(edge)) continue
+
+            if (createdLabels.add(edge.label)) {
+                val path = FlowPath.CfgEdge(edge.label, this.fir)
+                addAlternateFlow(path, buildAlternateFlow(path, builder).freeze())
             }
         }
     }
@@ -1307,11 +1338,4 @@ abstract class FirDataFlowAnalyzer(
 
     private fun MutableFlow.commitOperationStatement(statement: OperationStatement) =
         addAllStatements(logicSystem.approveOperationStatement(this, statement, removeApprovedOrImpossible = true))
-
-    private enum class FlowPathOperation {
-        PROPAGATE,
-        ADDITIONAL,
-        DIVERGE,
-        TERMINATE,
-    }
 }

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.resolve.dfa.cfg
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.contracts.description.*
 import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.hasExplicitBackingField
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
@@ -75,6 +76,7 @@ class ControlFlowGraphBuilder {
     private val catchBlocksInProgress: Stack<CatchClauseEnterNode> = stackOf()
     private val finallyEnterNodes: Stack<FinallyBlockEnterNode> = stackOf()
     private val finallyBlocksInProgress: Stack<FinallyBlockEnterNode> = stackOf()
+    private val finallyBlocksInProgressSet = mutableSetOf<FirElement>()
 
     private val exitSafeCallNodes: Stack<ExitSafeCallNode> = stackOf()
     private val exitElvisExpressionNodes: Stack<ElvisExitNode> = stackOf()
@@ -84,6 +86,10 @@ class ControlFlowGraphBuilder {
 
     // ----------------------------------- Public API -----------------------------------
 
+    fun withinFinallyBlock(element: FirElement): Boolean {
+        return finallyBlocksInProgressSet.contains(element)
+    }
+
     fun returnExpressionsOfAnonymousFunction(function: FirAnonymousFunction): Collection<FirAnonymousFunctionReturnExpressionInfo>? {
         val exitNode = function.controlFlowGraphReference?.controlFlowGraph?.exitNode ?: return null
 
@@ -964,6 +970,7 @@ class ControlFlowGraphBuilder {
         return finallyEnterNodes.pop().also {
             lastNodes.push(it)
             finallyBlocksInProgress.push(it)
+            finallyBlocksInProgressSet.add(it.fir)
         }
     }
 
@@ -1030,6 +1037,8 @@ class ControlFlowGraphBuilder {
         val node = tryExitNodes.pop()
         if (node.fir.finallyBlock != null) {
             val enterFinallyNode = finallyBlocksInProgress.pop()
+            finallyBlocksInProgressSet.remove(enterFinallyNode.fir)
+
             /**
              * If it appears that after completion try main expression returns nothing and try has finally block,
              *   we should make edge from finally exist to try exit a dead (and it may be not dead originally

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

@@ -170,22 +170,17 @@ sealed class CFGNode<out E : FirElement>(val owner: ControlFlowGraph, val level:
 
 val CFGNode<*>.firstPreviousNode: CFGNode<*> get() = previousNodes[0]
 val CFGNode<*>.lastPreviousNode: CFGNode<*> get() = previousNodes.last()
-val CFGNode<*>.previousDfaNodes: Sequence<Pair<Edge, CFGNode<*>>>
+fun CFGNode<*>.usedInDfa(edge: Edge) = if (isDead) edge.kind.usedInDeadDfa else edge.kind.usedInDfa
-    get() = previousNodes.asSequence()
+val CFGNode<*>.previousLiveNodes: List<CFGNode<*>>
-        .map { edgeFrom(it) to it }
-        .filter { (edge, _) -> if (isDead) edge.kind.usedInDeadDfa else edge.kind.usedInDfa }
-val CFGNode<*>.previousLiveNodes: Sequence<CFGNode<*>>
     get() = when  {
-        this.isDead -> previousNodes.asSequence()
+        this.isDead -> previousNodes
-        else -> previousNodes.asSequence().mapNotNull { it.takeIf { !it.isDead } }
+        else -> previousNodes.filter { !it.isDead }
     }
 
 interface EnterNodeMarker
 interface ExitNodeMarker
 interface GraphEnterNodeMarker : EnterNodeMarker
 interface GraphExitNodeMarker : ExitNodeMarker
-interface AlternateFlowStartMarker
-interface AlternateFlowEndMarker
 
 // ----------------------------------- EnterNode for declaration with CFG -----------------------------------
 
@@ -565,13 +560,13 @@ class CatchClauseExitNode(owner: ControlFlowGraph, override val fir: FirCatch, l
     }
 }
 class FinallyBlockEnterNode(owner: ControlFlowGraph, override val fir: FirTryExpression, level: Int) : CFGNode<FirTryExpression>(owner, level),
-    EnterNodeMarker, AlternateFlowStartMarker {
+    EnterNodeMarker {
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitFinallyBlockEnterNode(this, data)
     }
 }
 class FinallyBlockExitNode(owner: ControlFlowGraph, override val fir: FirTryExpression, level: Int) : CFGNode<FirTryExpression>(owner, level),
-    ExitNodeMarker, AlternateFlowEndMarker {
+    ExitNodeMarker {
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitFinallyBlockExitNode(this, data)
     }