FIR DFA: attach PersistentFlow directly to nodes

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

@@ -87,7 +87,7 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
             }
         }
 
-        var flow = dataFlowInfo.flowOnNodes.getValue(node)
+        var flow = node.flow
         val operation = effect.value.toOperation()
         if (operation != null) {
             if (resultExpression is FirConstExpression<*>) {

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

@@ -39,12 +39,9 @@ import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 class DataFlowAnalyzerContext(
     val graphBuilder: ControlFlowGraphBuilder,
     variableStorage: VariableStorageImpl,
-    flowOnNodes: MutableMap<CFGNode<*>, PersistentFlow>,
     val preliminaryLoopVisitor: PreliminaryLoopVisitor,
     val variablesClearedBeforeLoop: Stack<List<RealVariable>>,
 ) {
-    var flowOnNodes = flowOnNodes
-        private set
     var variableStorage = variableStorage
         private set
 
@@ -58,10 +55,7 @@ class DataFlowAnalyzerContext(
 
     fun reset() {
         graphBuilder.reset()
-
         variableStorage = variableStorage.clear()
-        flowOnNodes = mutableMapOf()
-
         preliminaryLoopVisitor.resetState()
         variablesClearedBeforeLoop.reset()
         firLocalVariableAssignmentAnalyzer = null
@@ -71,7 +65,7 @@ class DataFlowAnalyzerContext(
         fun empty(session: FirSession): DataFlowAnalyzerContext =
             DataFlowAnalyzerContext(
                 ControlFlowGraphBuilder(), VariableStorageImpl(session),
-                mutableMapOf(), PreliminaryLoopVisitor(), stackOf()
+                PreliminaryLoopVisitor(), stackOf()
             )
     }
 }
@@ -210,15 +204,13 @@ abstract class FirDataFlowAnalyzer(
                 variableStorage.removeRealVariable(valueParameter.symbol)
             }
         }
-        val variableStorage = variableStorage
-        val flowOnNodes = context.flowOnNodes
-
+        val info = DataFlowInfo(variableStorage)
         if (graphBuilder.isTopLevel()) {
             context.reset()
         } else {
             resetReceivers()
         }
-        return FirControlFlowGraphReferenceImpl(graph, DataFlowInfo(variableStorage, flowOnNodes))
+        return FirControlFlowGraphReferenceImpl(graph, info)
     }
 
     // ----------------------------------- Anonymous function -----------------------------------
@@ -253,6 +245,7 @@ abstract class FirDataFlowAnalyzer(
         }
         functionExitNode.mergeIncomingFlow()
         if (postponedLambdaExitNode != null) {
+            // Data flow changed, need to recompute. TODO: this violates the "flow computed only once" principle.
             postponedLambdaExitNode.mergeIncomingFlow()
         } else {
             resetReceivers()
@@ -1208,11 +1201,6 @@ abstract class FirDataFlowAnalyzer(
         }
     }
 
-    private val CFGNode<*>.flow: PersistentFlow
-        get() = context.flowOnNodes.getValue(this.origin)
-
-    private val CFGNode<*>.origin: CFGNode<*> get() = if (this is StubNode) firstPreviousNode else this
-
     private val CFGNode<*>.livePreviousFlows: List<PersistentFlow>
         get() = previousNodes.mapNotNull { it.takeIf { this.isDead || !it.isDead }?.flow }
 
@@ -1242,13 +1230,14 @@ abstract class FirDataFlowAnalyzer(
         return result
     }
 
+    @OptIn(CfgInternals::class)
     private fun UnionFunctionCallArgumentsNode.unionFlowFromArguments() {
-        context.flowOnNodes[this] = logicSystem.joinFlow(previousNodes.map { it.flow }, union = true).freeze()
+        flow = logicSystem.joinFlow(previousNodes.map { it.flow }, union = true).freeze()
     }
 
+    @OptIn(CfgInternals::class)
     private fun CFGNode<*>.setFlow(builder: MutableFlow): PersistentFlow {
-        val flow = builder.freeze()
-        context.flowOnNodes[this] = flow
+        val flow = builder.freeze().also { this.flow = it }
         if (currentReceiverState === builder) {
             currentReceiverState = flow
         }

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

@@ -25,7 +25,7 @@ class FirControlFlowGraphReferenceImpl(
     }
 }
 
-class DataFlowInfo(val variableStorage: VariableStorage, val flowOnNodes: Map<CFGNode<*>, PersistentFlow>)
+class DataFlowInfo(val variableStorage: VariableStorage)
 
 val FirControlFlowGraphReference.controlFlowGraph: ControlFlowGraph?
     get() = (this as? FirControlFlowGraphReferenceImpl)?.controlFlowGraph

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.resolve.dfa.PersistentFlow
 import org.jetbrains.kotlin.fir.resolve.dfa.controlFlowGraph
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
@@ -125,6 +126,12 @@ sealed class CFGNode<out E : FirElement>(val owner: ControlFlowGraph, val level:
     var isDead: Boolean = false
         protected set
 
+    private var _flow: PersistentFlow? = null
+    open var flow: PersistentFlow
+        get() = _flow ?: throw IllegalStateException("flow for $this not initialized - traversing nodes in wrong order?")
+        @CfgInternals
+        set(value) { _flow = value } // TODO: forbid reassignment
+
     @CfgInternals
     fun updateDeadStatus() {
         isDead = incomingEdges.size == previousNodes.size && incomingEdges.values.all {
@@ -747,6 +754,11 @@ class StubNode(owner: ControlFlowGraph, level: Int, id: Int) : CFGNode<FirStub>(
 
     override val fir: FirStub get() = FirStub
 
+    override var flow: PersistentFlow
+        get() = firstPreviousNode.flow
+        @CfgInternals
+        set(_) = throw IllegalStateException("can't set flow for stub node")
+
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitStubNode(this, data)
     }