[FIR] Add VariableStorage and flow on nodes into CFG reference

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

@@ -7,11 +7,15 @@ package org.jetbrains.kotlin.fir.resolve.dfa
 
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraph
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
-class FirControlFlowGraphReferenceImpl(val controlFlowGraph: ControlFlowGraph) : FirControlFlowGraphReference() {
+class FirControlFlowGraphReferenceImpl(
+    val controlFlowGraph: ControlFlowGraph,
+    val dataFlowInfo: DataFlowInfo? = null
+) : FirControlFlowGraphReference() {
     override val source: FirSourceElement? get() = null
 
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {}
@@ -21,5 +25,10 @@ class FirControlFlowGraphReferenceImpl(val controlFlowGraph: ControlFlowGraph) :
     }
 }
 
+class DataFlowInfo(val variableStorage: VariableStorage, val flowOnNodes: Map<CFGNode<*>, Flow>)
+
 val FirControlFlowGraphReference.controlFlowGraph: ControlFlowGraph?
-    get() = (this as? FirControlFlowGraphReferenceImpl)?.controlFlowGraph
+    get() = (this as? FirControlFlowGraphReferenceImpl)?.controlFlowGraph
+
+val FirControlFlowGraphReference.dataFlowInfo: DataFlowInfo?
+    get() = (this as? FirControlFlowGraphReferenceImpl)?.dataFlowInfo

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.contracts.description.ConeConstantReference
 import org.jetbrains.kotlin.fir.contracts.description.ConeReturnsEffectDeclaration
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.PersistentImplicitReceiverStack
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
@@ -33,15 +34,21 @@ import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 class DataFlowAnalyzerContext<FLOW : Flow>(
     val graphBuilder: ControlFlowGraphBuilder,
-    val variableStorage: VariableStorage,
-    val flowOnNodes: MutableMap<CFGNode<*>, FLOW>,
+    variableStorage: VariableStorage,
+    flowOnNodes: MutableMap<CFGNode<*>, FLOW>,
     val variablesForWhenConditions: MutableMap<WhenBranchConditionExitNode, DataFlowVariable>
 ) {
+    var flowOnNodes = flowOnNodes
+        private set
+    var variableStorage = variableStorage
+        private set
+
     fun reset() {
         graphBuilder.reset()
-        variableStorage.reset()
-        flowOnNodes.clear()
         variablesForWhenConditions.clear()
+
+        variableStorage = variableStorage.clear()
+        flowOnNodes = mutableMapOf()
     }
 
     companion object {
@@ -160,7 +167,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         functionEnterNode.mergeIncomingFlow(shouldForkFlow = previousNode != null)
     }
 
-    fun exitFunction(function: FirFunction<*>): ControlFlowGraph {
+    fun exitFunction(function: FirFunction<*>): FirControlFlowGraphReference {
         if (function is FirAnonymousFunction) {
             return exitAnonymousFunction(function)
         }
@@ -171,12 +178,13 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                 variableStorage.removeRealVariable(valueParameter.symbol)
             }
         }
+        val variableStorage = variableStorage
+        val flowOnNodes = context.flowOnNodes
+
         if (graphBuilder.isTopLevel()) {
-            context.flowOnNodes.clear()
-            variableStorage.reset()
-            graphBuilder.reset()
+            context.reset()
         }
-        return graph
+        return FirControlFlowGraphReferenceImpl(graph, DataFlowInfo(variableStorage, flowOnNodes))
     }
 
     // ----------------------------------- Anonymous function -----------------------------------
@@ -188,12 +196,12 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         functionEnterNode.mergeIncomingFlow()
     }
 
-    private fun exitAnonymousFunction(anonymousFunction: FirAnonymousFunction): ControlFlowGraph {
+    private fun exitAnonymousFunction(anonymousFunction: FirAnonymousFunction): FirControlFlowGraphReference {
         val (functionExitNode, postponedLambdaExitNode, graph) = graphBuilder.exitAnonymousFunction(anonymousFunction)
         // TODO: questionable
         postponedLambdaExitNode?.mergeIncomingFlow()
         functionExitNode.mergeIncomingFlow()
-        return graph
+        return FirControlFlowGraphReferenceImpl(graph)
     }
 
     fun visitPostponedAnonymousFunction(anonymousFunction: FirAnonymousFunction) {

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

@@ -31,6 +31,8 @@ class VariableStorage(private val session: FirSession) {
     private val realVariables: MutableMap<Identifier, RealVariable> = HashMap()
     private val syntheticVariables: MutableMap<FirElement, SyntheticVariable> = HashMap()
 
+    fun clear(): VariableStorage = VariableStorage(session)
+
     fun getOrCreateRealVariableWithoutUnwrappingAlias(flow: Flow, symbol: AbstractFirBasedSymbol<*>, fir: FirElement): RealVariable {
         val realFir = fir.unwrapElement()
         val identifier = getIdentifierBySymbol(flow, symbol, realFir)
@@ -139,12 +141,6 @@ class VariableStorage(private val session: FirSession) {
         syntheticVariables.remove(variable.fir)
     }
 
-    fun reset() {
-        counter = 0
-        realVariables.clear()
-        syntheticVariables.clear()
-    }
-
     @OptIn(ExperimentalContracts::class)
     fun AbstractFirBasedSymbol<*>?.isStable(originalFir: FirElement): Boolean {
         contract {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -521,8 +521,8 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             transformDeclarationContent(function, data).also {
                 if (functionIsNotAnalyzed) {
                     val result = it.single as FirFunction<*>
-                    val controlFlowGraph = dataFlowAnalyzer.exitFunction(result)
-                    result.replaceControlFlowGraphReference(FirControlFlowGraphReferenceImpl(controlFlowGraph))
+                    val controlFlowGraphReference = dataFlowAnalyzer.exitFunction(result)
+                    result.replaceControlFlowGraphReference(controlFlowGraphReference)
                 }
             } as CompositeTransformResult<FirStatement>
         }
@@ -593,8 +593,8 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                 }
             }
 
-            val graph = dataFlowAnalyzer.exitFunction(constructor)
-            constructor.replaceControlFlowGraphReference(FirControlFlowGraphReferenceImpl(graph))
+            val controlFlowGraphReference = dataFlowAnalyzer.exitFunction(constructor)
+            constructor.replaceControlFlowGraphReference(controlFlowGraphReference)
             constructor.compose()
         }
     }