From dc57d307f3ba5993dc721118cf9994fcf435a979 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Tue, 19 May 2020 12:01:49 +0300
Subject: [PATCH] [FIR] Clean data flow context before body resolve of file

---
 .../kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt  |  7 +++++++
 .../org/jetbrains/kotlin/fir/resolve/dfa/Stack.kt  | 14 ++++++++++++++
 .../fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt |  2 ++
 .../resolve/FirAbstractBodyResolveTransformer.kt   |  4 ++++
 .../body/resolve/FirBodyResolveTransformer.kt      |  1 +
 5 files changed, 28 insertions(+)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
index b975cd5d84e..148918cd292 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
@@ -36,6 +36,13 @@ class DataFlowAnalyzerContext<FLOW : Flow>(
     val flowOnNodes: MutableMap<CFGNode<*>, FLOW>,
     val variablesForWhenConditions: MutableMap<WhenBranchConditionExitNode, DataFlowVariable>
 ) {
+    fun reset() {
+        graphBuilder.reset()
+        variableStorage.reset()
+        flowOnNodes.clear()
+        variablesForWhenConditions.clear()
+    }
+
     companion object {
         fun <FLOW : Flow> empty(session: FirSession) =
             DataFlowAnalyzerContext<FLOW>(
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/Stack.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/Stack.kt
index ffe981fb57e..43dfdfd8005 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/Stack.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/Stack.kt
@@ -15,6 +15,7 @@ abstract class Stack<T> {
     abstract fun top(): T
     abstract fun pop(): T
     abstract fun push(value: T)
+    abstract fun reset()
 }
 
 fun <T> stackOf(vararg values: T): Stack<T> = StackImpl(*values)
@@ -33,6 +34,9 @@ private class StackImpl<T>(vararg values: T) : Stack<T>() {
     }
 
     override val size: Int get() = stack.size
+    override fun reset() {
+        stack.clear()
+    }
 }
 
 class NodeStorage<T : FirElement, N : CFGNode<T>> : Stack<N>(){
@@ -55,6 +59,11 @@ class NodeStorage<T : FirElement, N : CFGNode<T>> : Stack<N>(){
     operator fun get(key: T): N? {
         return map[key]
     }
+
+    override fun reset() {
+        stack.reset()
+        map.clear()
+    }
 }
 
 class SymbolBasedNodeStorage<T, N : CFGNode<T>> : Stack<N>() where T : FirElement {
@@ -77,4 +86,9 @@ class SymbolBasedNodeStorage<T, N : CFGNode<T>> : Stack<N>() where T : FirElemen
     operator fun get(key: FirBasedSymbol<*>): N? {
         return map[key]
     }
+
+    override fun reset() {
+        stack.reset()
+        map.clear()
+    }
 }
\ No newline at end of file
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt
index 1518a0e858a..bf4547b563e 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt
@@ -941,6 +941,8 @@ class ControlFlowGraphBuilder {
     fun reset() {
         exitsOfAnonymousFunctions.clear()
         exitsFromCompletedPostponedAnonymousFunctions.clear()
+        lexicalScopes.reset()
+        lexicalScopes.push(stackOf())
     }
 
     fun dropSubgraphFromCall(call: FirFunctionCall) {
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformer.kt
index 123d60aea04..bff1bedf5fb 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformer.kt
@@ -229,6 +229,10 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
             towerDataContextForAnonymousFunctions.clear()
         }
 
+        fun cleanDataFlowContext() {
+            dataFlowAnalyzerContext.reset()
+        }
+
         private inline fun updateLastScope(transform: FirLocalScope.() -> FirLocalScope) {
             val lastScope = towerDataContext.localScopes.lastOrNull() ?: return
             replaceTowerDataContext(towerDataContext.setLastLocalScope(lastScope.transform()))
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt
index 86f475b4400..72059b7a05d 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt
@@ -46,6 +46,7 @@ open class FirBodyResolveTransformer(
     override fun transformFile(file: FirFile, data: ResolutionMode): CompositeTransformResult<FirFile> {
         checkSessionConsistency(file)
         context.cleanContextForAnonymousFunction()
+        context.cleanDataFlowContext()
         @OptIn(PrivateForInline::class)
         context.file = file
         packageFqName = file.packageFqName