FIR CFA: require path awareness in fixed point iteration

It really doesn't make sense to collect control flow data without
checking paths, since that'll produce incorrect results for try-catch.

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/FirCallsEffectAnalyzer.kt

@@ -222,9 +222,6 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
 
         override val constructor: (PersistentMap<FirBasedSymbol<*>, EventOccurrencesRange>) -> LambdaInvocationInfo =
             ::LambdaInvocationInfo
-
-        override val empty: () -> LambdaInvocationInfo =
-            ::EMPTY
     }
 
     class PathAwareLambdaInvocationInfo(
@@ -236,9 +233,6 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
 
         override val constructor: (PersistentMap<EdgeLabel, LambdaInvocationInfo>) -> PathAwareLambdaInvocationInfo =
             ::PathAwareLambdaInvocationInfo
-
-        override val empty: () -> PathAwareLambdaInvocationInfo =
-            ::EMPTY
     }
 
     private class InvocationDataCollector(
@@ -249,7 +243,7 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
 
         override fun visitFunctionCallNode(
             node: FunctionCallNode,
-            data: Collection<Pair<EdgeLabel, PathAwareLambdaInvocationInfo>>
+            data: PathAwareLambdaInvocationInfo
         ): PathAwareLambdaInvocationInfo {
             var dataForNode = visitUnionNode(node, data)
 

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/util/CfgTraverser.kt

@@ -33,10 +33,10 @@ fun ControlFlowGraph.traverse(
 
 // ---------------------- Path-sensitive data collection -----------------------
 
-fun <I> ControlFlowGraph.collectDataForNode(
+fun <I : PathAwareControlFlowInfo<I, *>> ControlFlowGraph.collectDataForNode(
     direction: TraverseDirection,
     initialInfo: I,
-    visitor: ControlFlowGraphVisitor<I, Collection<Pair<EdgeLabel, I>>>,
+    visitor: PathAwareControlFlowGraphVisitor<I>,
     visitSubGraphs: Boolean = true
 ): Map<CFGNode<*>, I> {
     val nodeMap = LinkedHashMap<CFGNode<*>, I>()
@@ -51,10 +51,10 @@ fun <I> ControlFlowGraph.collectDataForNode(
     return nodeMap
 }
 
-private fun <I> ControlFlowGraph.collectDataForNodeInternal(
+private fun <I : PathAwareControlFlowInfo<I, *>> ControlFlowGraph.collectDataForNodeInternal(
     direction: TraverseDirection,
     initialInfo: I,
-    visitor: ControlFlowGraphVisitor<I, Collection<Pair<EdgeLabel, I>>>,
+    visitor: PathAwareControlFlowGraphVisitor<I>,
     nodeMap: MutableMap<CFGNode<*>, I>,
     visitSubGraphs: Boolean = true
 ): Boolean {
@@ -68,18 +68,21 @@ private fun <I> ControlFlowGraph.collectDataForNodeInternal(
             TraverseDirection.Forward -> node.previousCfgNodes
             TraverseDirection.Backward -> node.followingCfgNodes
         }
-        // One noticeable different against the path-unaware version is, here, we pair the control-flow info with the label.
         val previousData =
             previousNodes.mapNotNull {
                 val k = when (direction) {
-                    TraverseDirection.Forward -> node.edgeFrom(it).label
-                    TraverseDirection.Backward -> node.edgeTo(it).label
+                    TraverseDirection.Forward -> node.edgeFrom(it)
+                    TraverseDirection.Backward -> node.edgeTo(it)
                 }
                 val v = nodeMap[it] ?: return@mapNotNull null
-                k to v
+                visitor.visitEdge(it, node, k, v)
             }
+        val reduced = if (node is UnionNodeMarker)
+            previousData.reduceOrNull { a, b -> a.plus(b) }
+        else
+            previousData.reduceOrNull { a, b -> a.merge(b) }
         val data = nodeMap[node]
-        val newData = node.accept(visitor, previousData)
+        val newData = node.accept(visitor, reduced ?: visitor.emptyInfo)
         val hasChanged = newData != data
         changed = changed or hasChanged
         if (hasChanged) {

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/util/ControlFlowInfo.kt

@@ -13,8 +13,6 @@ abstract class ControlFlowInfo<S : ControlFlowInfo<S, K, V>, K : Any, V : Any> p
 
     protected abstract val constructor: (PersistentMap<K, V>) -> S
 
-    protected abstract val empty: () -> S
-
     override fun equals(other: Any?): Boolean {
         return map == (other as? ControlFlowInfo<*, *, *>)?.map
     }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/util/PathAwareControlFlowGraphVisitor.kt

@@ -7,22 +7,16 @@ package org.jetbrains.kotlin.fir.analysis.cfa.util
 
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraphVisitor
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.EdgeLabel
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.Edge
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.UnionNodeMarker
-import org.jetbrains.kotlin.utils.addToStdlib.foldMap
 
-abstract class PathAwareControlFlowGraphVisitor<P : PathAwareControlFlowInfo<P, *>>
-    : ControlFlowGraphVisitor<P, Collection<Pair<EdgeLabel, P>>>() {
+abstract class PathAwareControlFlowGraphVisitor<P : PathAwareControlFlowInfo<P, *>> : ControlFlowGraphVisitor<P, P>() {
+    abstract val emptyInfo: P
 
-    protected abstract val emptyInfo: P
+    open fun visitEdge(from: CFGNode<*>, to: CFGNode<*>, metadata: Edge, data: P): P =
+        data.applyLabel(to, metadata.label) ?: emptyInfo
 
-    override fun visitNode(node: CFGNode<*>, data: Collection<Pair<EdgeLabel, P>>): P {
-        if (data.isEmpty()) return emptyInfo
-        return data.foldMap({ (label, info) -> info.applyLabel(node, label) }) { a, b -> a.merge(b) }
-    }
+    override fun visitNode(node: CFGNode<*>, data: P): P = data
 
-    override fun <T> visitUnionNode(node: T, data: Collection<Pair<EdgeLabel, P>>): P where T : CFGNode<*>, T : UnionNodeMarker {
-        if (data.isEmpty()) return emptyInfo
-        return data.foldMap({ (label, info) -> info.applyLabel(node, label) }) { a, b -> a.plus(b) }
-    }
+    override fun <T> visitUnionNode(node: T, data: P): P where T : CFGNode<*>, T : UnionNodeMarker = data
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/util/PathAwareControlFlowInfo.kt

@@ -16,7 +16,7 @@ abstract class PathAwareControlFlowInfo<P : PathAwareControlFlowInfo<P, S>, S :
     internal val infoAtNormalPath: S
         get() = map.getValue(NormalPath)
 
-    fun applyLabel(node: CFGNode<*>, label: EdgeLabel): P {
+    fun applyLabel(node: CFGNode<*>, label: EdgeLabel): P? {
         if (label.isNormal) {
             // Special case: when we exit the try expression, null label means a normal path.
             // Filter out any info bound to non-null label
@@ -28,7 +28,7 @@ abstract class PathAwareControlFlowInfo<P : PathAwareControlFlowInfo<P, S>, S :
                     constructor(persistentMapOf(NormalPath to infoAtNormalPath))
                 } else {
                     /* This means no info for normal path. */
-                    empty()
+                    null
                 }
             }
             // In general, null label means no additional path info, hence return `this` as-is.
@@ -58,7 +58,7 @@ abstract class PathAwareControlFlowInfo<P : PathAwareControlFlowInfo<P, S>, S :
                 }
             } else {
                 /* This means no info for the specific label. */
-                empty()
+                null
             }
         } else {
             // { |-> ... }    // empty path info

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/util/PropertyInitializationInfo.kt

@@ -49,9 +49,6 @@ class PropertyInitializationInfo(
 
     override val constructor: (PersistentMap<FirPropertySymbol, EventOccurrencesRange>) -> PropertyInitializationInfo =
         ::PropertyInitializationInfo
-
-    override val empty: () -> PropertyInitializationInfo =
-        ::EMPTY
 }
 
 class PathAwarePropertyInitializationInfo(
@@ -63,7 +60,4 @@ class PathAwarePropertyInitializationInfo(
 
     override val constructor: (PersistentMap<EdgeLabel, PropertyInitializationInfo>) -> PathAwarePropertyInitializationInfo =
         ::PathAwarePropertyInitializationInfo
-
-    override val empty: () -> PathAwarePropertyInitializationInfo =
-        ::EMPTY
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/util/PropertyInitializationInfoCollector.kt

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.fir.analysis.cfa.util
 import kotlinx.collections.immutable.PersistentMap
 import kotlinx.collections.immutable.persistentMapOf
 import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.references.toResolvedPropertySymbol
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
@@ -32,7 +31,7 @@ class PropertyInitializationInfoCollector(
 
     override fun visitVariableAssignmentNode(
         node: VariableAssignmentNode,
-        data: Collection<Pair<EdgeLabel, PathAwarePropertyInitializationInfo>>
+        data: PathAwarePropertyInitializationInfo
     ): PathAwarePropertyInitializationInfo {
         val dataForNode = visitNode(node, data)
         val symbol = node.fir.calleeReference.toResolvedPropertySymbol() ?: return dataForNode
@@ -45,7 +44,7 @@ class PropertyInitializationInfoCollector(
 
     override fun visitVariableDeclarationNode(
         node: VariableDeclarationNode,
-        data: Collection<Pair<EdgeLabel, PathAwarePropertyInitializationInfo>>
+        data: PathAwarePropertyInitializationInfo
     ): PathAwarePropertyInitializationInfo {
         val dataForNode = visitNode(node, data)
         return processVariableWithAssignment(
@@ -78,72 +77,35 @@ class PropertyInitializationInfoCollector(
     // Data flows of declared/assigned variables in loops
     // --------------------------------------------------
 
-    private fun enterCapturingStatement(statement: FirStatement): Set<FirPropertySymbol> =
-        declaredVariableCollector.enterCapturingStatement(statement)
-
-    private fun exitCapturingStatement(statement: FirStatement) {
-        declaredVariableCollector.exitCapturingStatement(statement)
-    }
-
-    // A merge point for a loop with `continue`
-    override fun visitLoopEnterNode(
-        node: LoopEnterNode,
-        data: Collection<Pair<EdgeLabel, PathAwarePropertyInitializationInfo>>
+    override fun visitEdge(
+        from: CFGNode<*>,
+        to: CFGNode<*>,
+        metadata: Edge,
+        data: PathAwarePropertyInitializationInfo
     ): PathAwarePropertyInitializationInfo {
-        val declaredVariableSymbolsInLoop = enterCapturingStatement(node.fir)
-        if (declaredVariableSymbolsInLoop.isEmpty())
-            return visitNode(node, data)
-
-        return filterDeclaredVariableSymbolsInCapturedScope(node, declaredVariableSymbolsInLoop, data)
-    }
-
-    // A merge point for while loop
-    override fun visitLoopConditionEnterNode(
-        node: LoopConditionEnterNode,
-        data: Collection<Pair<EdgeLabel, PathAwarePropertyInitializationInfo>>
-    ): PathAwarePropertyInitializationInfo {
-        val declaredVariableSymbolsInLoop = declaredVariableCollector.declaredVariablesPerElement[node.loop]
-        if (declaredVariableSymbolsInLoop.isEmpty())
-            return visitNode(node, data)
-
-        return filterDeclaredVariableSymbolsInCapturedScope(node, declaredVariableSymbolsInLoop, data)
-    }
-
-    // A merge point for do-while loop
-    override fun visitLoopBlockEnterNode(
-        node: LoopBlockEnterNode,
-        data: Collection<Pair<EdgeLabel, PathAwarePropertyInitializationInfo>>
-    ): PathAwarePropertyInitializationInfo {
-        val declaredVariableSymbolsInLoop = declaredVariableCollector.declaredVariablesPerElement[node.fir]
-        if (declaredVariableSymbolsInLoop.isEmpty())
-            return visitNode(node, data)
-
-        return filterDeclaredVariableSymbolsInCapturedScope(node, declaredVariableSymbolsInLoop, data)
-    }
-
-    private fun filterDeclaredVariableSymbolsInCapturedScope(
-        node: CFGNode<*>,
-        declaredVariableSymbolsInCapturedScope: Collection<FirPropertySymbol>,
-        data: Collection<Pair<EdgeLabel, PathAwarePropertyInitializationInfo>>
-    ): PathAwarePropertyInitializationInfo {
-        var filteredData = data
-        for (variableSymbol in declaredVariableSymbolsInCapturedScope) {
-            filteredData = filteredData.map { (label, pathAwareInfo) ->
-                label to if (label is LoopBackPath) {
-                    removeRange(pathAwareInfo, variableSymbol, ::PathAwarePropertyInitializationInfo)
-                } else {
-                    pathAwareInfo
-                }
-            }
+        val result = super.visitEdge(from, to, metadata, data)
+        if (metadata.label != LoopBackPath) return result
+        val declaredVariableSymbolsInCapturedScope = when (to) {
+            is LoopEnterNode -> declaredVariableCollector.declaredVariablesPerElement[to.fir]
+            is LoopBlockEnterNode -> declaredVariableCollector.declaredVariablesPerElement[to.fir]
+            is LoopConditionEnterNode -> declaredVariableCollector.declaredVariablesPerElement[to.loop]
+            else -> return result
         }
-        return visitNode(node, filteredData)
+        return declaredVariableSymbolsInCapturedScope.fold(data) { filteredData, variableSymbol ->
+            removeRange(filteredData, variableSymbol, ::PathAwarePropertyInitializationInfo)
+        }
+    }
+
+    override fun visitLoopEnterNode(node: LoopEnterNode, data: PathAwarePropertyInitializationInfo): PathAwarePropertyInitializationInfo {
+        declaredVariableCollector.enterCapturingStatement(node.fir)
+        return visitNode(node, data)
     }
 
     override fun visitLoopExitNode(
         node: LoopExitNode,
-        data: Collection<Pair<EdgeLabel, PathAwarePropertyInitializationInfo>>
+        data: PathAwarePropertyInitializationInfo
     ): PathAwarePropertyInitializationInfo {
-        exitCapturingStatement(node.fir)
+        declaredVariableCollector.exitCapturingStatement(node.fir)
         return visitNode(node, data)
     }
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/UnusedChecker.kt

@@ -136,9 +136,6 @@ object UnusedChecker : AbstractFirPropertyInitializationChecker() {
         override val constructor: (PersistentMap<FirPropertySymbol, VariableStatus>) -> VariableStatusInfo =
             ::VariableStatusInfo
 
-        override val empty: () -> VariableStatusInfo =
-            ::EMPTY
-
         override fun merge(other: VariableStatusInfo): VariableStatusInfo {
             var result = this
             for (symbol in keys.union(other.keys)) {
@@ -163,9 +160,6 @@ object UnusedChecker : AbstractFirPropertyInitializationChecker() {
 
         override val constructor: (PersistentMap<EdgeLabel, VariableStatusInfo>) -> PathAwareVariableStatusInfo =
             ::PathAwareVariableStatusInfo
-
-        override val empty: () -> PathAwareVariableStatusInfo =
-            ::EMPTY
     }
 
     private class ValueWritesWithoutReading(
@@ -184,19 +178,19 @@ object UnusedChecker : AbstractFirPropertyInitializationChecker() {
 
         override fun visitNode(
             node: CFGNode<*>,
-            data: Collection<Pair<EdgeLabel, PathAwareVariableStatusInfo>>
+            data: PathAwareVariableStatusInfo
         ): PathAwareVariableStatusInfo =
             super.visitNode(node, data).withAnnotationsFrom(node)
 
         override fun <T> visitUnionNode(
             node: T,
-            data: Collection<Pair<EdgeLabel, PathAwareVariableStatusInfo>>
+            data: PathAwareVariableStatusInfo
         ): PathAwareVariableStatusInfo where T : CFGNode<*>, T : UnionNodeMarker =
             super.visitUnionNode(node, data).withAnnotationsFrom(node)
 
         override fun visitVariableDeclarationNode(
             node: VariableDeclarationNode,
-            data: Collection<Pair<EdgeLabel, PathAwareVariableStatusInfo>>
+            data: PathAwareVariableStatusInfo
         ): PathAwareVariableStatusInfo {
             val dataForNode = visitNode(node, data)
             if (node.fir.source?.kind is KtFakeSourceElementKind) return dataForNode
@@ -228,7 +222,7 @@ object UnusedChecker : AbstractFirPropertyInitializationChecker() {
 
         override fun visitVariableAssignmentNode(
             node: VariableAssignmentNode,
-            data: Collection<Pair<EdgeLabel, PathAwareVariableStatusInfo>>
+            data: PathAwareVariableStatusInfo
         ): PathAwareVariableStatusInfo {
             val dataForNode = visitNode(node, data)
             val symbol = node.fir.lValue.toResolvedPropertySymbol() ?: return dataForNode
@@ -257,7 +251,7 @@ object UnusedChecker : AbstractFirPropertyInitializationChecker() {
 
         override fun visitQualifiedAccessNode(
             node: QualifiedAccessNode,
-            data: Collection<Pair<EdgeLabel, PathAwareVariableStatusInfo>>
+            data: PathAwareVariableStatusInfo
         ): PathAwareVariableStatusInfo {
             val dataForNode = visitNode(node, data)
             return visitQualifiedAccesses(dataForNode, node.fir)
@@ -293,7 +287,7 @@ object UnusedChecker : AbstractFirPropertyInitializationChecker() {
 
         override fun visitFunctionCallNode(
             node: FunctionCallNode,
-            data: Collection<Pair<EdgeLabel, PathAwareVariableStatusInfo>>
+            data: PathAwareVariableStatusInfo
         ): PathAwareVariableStatusInfo {
             val dataForNode = visitUnionNode(node, data)
             val reference = node.fir.calleeReference.resolved ?: return dataForNode