FIR DFA: attach PersistentFlow directly to nodes

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)
     }