FirTowerDataContext parts: split into regular and special

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/BodyResolveComponents.kt

@@ -8,8 +8,6 @@ package org.jetbrains.kotlin.fir.resolve
 import org.jetbrains.kotlin.fir.FirCallResolver
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataMode.*
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionStageRunner
 import org.jetbrains.kotlin.fir.resolve.dfa.FirDataFlowAnalyzer
 import org.jetbrains.kotlin.fir.resolve.inference.FirCallCompleter
@@ -17,10 +15,7 @@ import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSyntheticCallGenerator
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.scopes.FirScope
-import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
-import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.types.FirTypeRef
-import java.util.*
 
 data class SessionHolderImpl(override val session: FirSession, override val scopeSession: ScopeSession) : SessionHolder {
     companion object {
@@ -50,64 +45,6 @@ abstract class BodyResolveComponents : SessionHolder {
     abstract val outerClassManager: FirOuterClassManager
 }
 
-enum class FirTowerDataMode {
-    MEMBER_DECLARATION,
-    NESTED_CLASS,
-    COMPANION_OBJECT,
-    CONSTRUCTOR_HEADER,
-    ENUM_ENTRY,
-    SPECIAL,
-}
-
-class FirTowerDataContextsForClassParts(
-    forMemberDeclarations: FirTowerDataContext,
-    forNestedClasses: FirTowerDataContext? = null,
-    forCompanionObject: FirTowerDataContext? = null,
-    forConstructorHeaders: FirTowerDataContext? = null,
-    forEnumEntries: FirTowerDataContext? = null,
-    val primaryConstructorPureParametersScope: FirLocalScope? = null,
-    val primaryConstructorAllParametersScope: FirLocalScope? = null,
-) {
-    private val modeMap = EnumMap<FirTowerDataMode, FirTowerDataContext>(FirTowerDataMode::class.java)
-
-    init {
-        modeMap[MEMBER_DECLARATION] = forMemberDeclarations
-        modeMap[NESTED_CLASS] = forNestedClasses
-        modeMap[COMPANION_OBJECT] = forCompanionObject
-        modeMap[CONSTRUCTOR_HEADER] = forConstructorHeaders
-        modeMap[ENUM_ENTRY] = forEnumEntries
-    }
-
-    var mode: FirTowerDataMode = MEMBER_DECLARATION
-
-    val forMemberDeclaration: FirTowerDataContext get() = modeMap.getValue(MEMBER_DECLARATION)
-
-    val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext> = mutableMapOf()
-    val towerDataContextForCallableReferences: MutableMap<FirCallableReferenceAccess, FirTowerDataContext> = mutableMapOf()
-
-    var currentContext: FirTowerDataContext
-        get() = modeMap.getValue(mode)
-        set(value) {
-            modeMap[mode] = value
-        }
-
-    fun setAnonymousFunctionContextIfAny(symbol: FirAnonymousFunctionSymbol) {
-        val context = towerDataContextForAnonymousFunctions[symbol]
-        if (context != null) {
-            mode = SPECIAL
-            modeMap[SPECIAL] = context
-        }
-    }
-
-    fun setCallableReferenceContextIfAny(access: FirCallableReferenceAccess) {
-        val context = towerDataContextForCallableReferences[access]
-        if (context != null) {
-            mode = SPECIAL
-            modeMap[SPECIAL] = context
-        }
-    }
-}
-
 // --------------------------------------- Utils ---------------------------------------
 
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/FirRegularTowerDataContexts.kt

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.resolve
+
+import org.jetbrains.kotlin.fir.declarations.FirTowerDataContext
+import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
+import java.util.*
+
+enum class FirTowerDataMode {
+    MEMBER_DECLARATION,
+    NESTED_CLASS,
+    COMPANION_OBJECT,
+    CONSTRUCTOR_HEADER,
+    ENUM_ENTRY,
+    SPECIAL,
+}
+
+class FirRegularTowerDataContexts(
+    forMemberDeclarations: FirTowerDataContext,
+    forNestedClasses: FirTowerDataContext? = null,
+    forCompanionObject: FirTowerDataContext? = null,
+    forConstructorHeaders: FirTowerDataContext? = null,
+    forEnumEntries: FirTowerDataContext? = null,
+    val primaryConstructorPureParametersScope: FirLocalScope? = null,
+    val primaryConstructorAllParametersScope: FirLocalScope? = null,
+) {
+    private val modeMap = EnumMap<FirTowerDataMode, FirTowerDataContext>(FirTowerDataMode::class.java)
+
+    init {
+        modeMap[FirTowerDataMode.MEMBER_DECLARATION] = forMemberDeclarations
+        modeMap[FirTowerDataMode.NESTED_CLASS] = forNestedClasses
+        modeMap[FirTowerDataMode.COMPANION_OBJECT] = forCompanionObject
+        modeMap[FirTowerDataMode.CONSTRUCTOR_HEADER] = forConstructorHeaders
+        modeMap[FirTowerDataMode.ENUM_ENTRY] = forEnumEntries
+    }
+
+    var mode: FirTowerDataMode = FirTowerDataMode.MEMBER_DECLARATION
+
+    var currentContext: FirTowerDataContext?
+        get() = modeMap[mode]
+        set(value) {
+            modeMap[mode] = value
+        }
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/FirSpecialTowerDataContexts.kt

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.resolve
+
+import org.jetbrains.kotlin.fir.declarations.FirTowerDataContext
+import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
+import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
+
+class FirSpecialTowerDataContexts {
+    val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext> = mutableMapOf()
+    val towerDataContextForCallableReferences: MutableMap<FirCallableReferenceAccess, FirTowerDataContext> = mutableMapOf()
+    var currentContext: FirTowerDataContext? = null
+
+    fun setAnonymousFunctionContextIfAny(symbol: FirAnonymousFunctionSymbol): Boolean {
+        val context = towerDataContextForAnonymousFunctions[symbol]
+        if (context != null) {
+            currentContext = context
+            return true
+        }
+        return false
+    }
+
+    fun setCallableReferenceContextIfAny(access: FirCallableReferenceAccess): Boolean {
+        val context = towerDataContextForCallableReferences[access]
+        if (context != null) {
+            currentContext = context
+            return true
+        }
+        return false
+    }
+}

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

@@ -15,7 +15,6 @@ import org.jetbrains.kotlin.fir.declarations.utils.isCompanion
 import org.jetbrains.kotlin.fir.declarations.utils.isInner
 import org.jetbrains.kotlin.fir.declarations.primaryConstructorIfAny
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
-import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitExtensionReceiverValue
@@ -53,27 +52,31 @@ class BodyResolveContext(
     @set:PrivateForInline
     lateinit var file: FirFile
 
-    var towerDataContextsForClassParts = FirTowerDataContextsForClassParts(forMemberDeclarations = FirTowerDataContext())
+    var regularTowerDataContexts = FirRegularTowerDataContexts(forMemberDeclarations = FirTowerDataContext())
         private set
 
+    val specialTowerDataContexts = FirSpecialTowerDataContexts()
+
     val towerDataContext: FirTowerDataContext
-        get() = towerDataContextsForClassParts.currentContext
+        get() = regularTowerDataContexts.currentContext
+            ?: specialTowerDataContexts.currentContext
+            ?: throw AssertionError("Not current data context found, towerDataMode = $towerDataMode")
 
     @OptIn(PrivateForInline::class)
     var towerDataMode: FirTowerDataMode
-        get() = towerDataContextsForClassParts.mode
+        get() = regularTowerDataContexts.mode
         set(value) {
-            towerDataContextsForClassParts.mode = value
+            regularTowerDataContexts.mode = value
         }
 
     val implicitReceiverStack: ImplicitReceiverStack
         get() = towerDataContext.implicitReceiverStack
 
     private val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext>
-        get() = towerDataContextsForClassParts.towerDataContextForAnonymousFunctions
+        get() = specialTowerDataContexts.towerDataContextForAnonymousFunctions
 
     private val towerDataContextForCallableReferences: MutableMap<FirCallableReferenceAccess, FirTowerDataContext>
-        get() = towerDataContextsForClassParts.towerDataContextForCallableReferences
+        get() = specialTowerDataContexts.towerDataContextForCallableReferences
 
     @set:PrivateForInline
     var containers: ArrayDeque<FirDeclaration> = ArrayDeque()
@@ -94,13 +97,13 @@ class BodyResolveContext(
     val topClassDeclaration: FirRegularClass?
         get() = containingClassDeclarations.lastOrNull()
 
-    private inline fun <T> withNewTowerDataForClassParts(newContexts: FirTowerDataContextsForClassParts, f: () -> T): T {
+    private inline fun <T> withNewTowerDataForClass(newContexts: FirRegularTowerDataContexts, f: () -> T): T {
-        val old = towerDataContextsForClassParts
+        val old = regularTowerDataContexts
-        towerDataContextsForClassParts = newContexts
+        regularTowerDataContexts = newContexts
         return try {
             f()
         } finally {
-            towerDataContextsForClassParts = old
+            regularTowerDataContexts = old
         }
     }
 
@@ -175,7 +178,11 @@ class BodyResolveContext(
 
     @PrivateForInline
     fun replaceTowerDataContext(newContext: FirTowerDataContext) {
-        towerDataContextsForClassParts.currentContext = newContext
+        if (towerDataMode == FirTowerDataMode.SPECIAL) {
+            specialTowerDataContexts.currentContext = newContext
+        } else {
+            regularTowerDataContexts.currentContext = newContext
+        }
     }
 
     @PrivateForInline
@@ -280,10 +287,10 @@ class BodyResolveContext(
     // ANALYSIS PUBLIC API
 
     fun getPrimaryConstructorPureParametersScope(): FirLocalScope? =
-        towerDataContextsForClassParts.primaryConstructorPureParametersScope
+        regularTowerDataContexts.primaryConstructorPureParametersScope
 
     fun getPrimaryConstructorAllParametersScope(): FirLocalScope? =
-        towerDataContextsForClassParts.primaryConstructorAllParametersScope
+        regularTowerDataContexts.primaryConstructorAllParametersScope
 
     @OptIn(PrivateForInline::class)
     fun storeClassIfNotNested(klass: FirRegularClass, session: FirSession) {
@@ -310,7 +317,9 @@ class BodyResolveContext(
     @OptIn(PrivateForInline::class)
     inline fun <T> withAnonymousFunctionTowerDataContext(symbol: FirAnonymousFunctionSymbol, f: () -> T): T {
         return withTowerModeCleanup {
-            towerDataContextsForClassParts.setAnonymousFunctionContextIfAny(symbol)
+            if (specialTowerDataContexts.setAnonymousFunctionContextIfAny(symbol)) {
+                regularTowerDataContexts.mode = FirTowerDataMode.SPECIAL
+            }
             f()
         }
     }
@@ -318,7 +327,9 @@ class BodyResolveContext(
     @OptIn(PrivateForInline::class)
     inline fun <T> withCallableReferenceTowerDataContext(access: FirCallableReferenceAccess, f: () -> T): T {
         return withTowerModeCleanup {
-            towerDataContextsForClassParts.setCallableReferenceContextIfAny(access)
+            if (specialTowerDataContexts.setCallableReferenceContextIfAny(access)) {
+                regularTowerDataContexts.mode = FirTowerDataMode.SPECIAL
+            }
             f()
         }
     }
@@ -460,7 +471,7 @@ class BodyResolveContext(
                 null to null
             }
 
-        val newContexts = FirTowerDataContextsForClassParts(
+        val newContexts = FirRegularTowerDataContexts(
             forMembersResolution,
             newTowerDataContextForStaticNestedClasses,
             statics,
@@ -470,7 +481,7 @@ class BodyResolveContext(
             primaryConstructorAllParametersScope
         )
 
-        return withNewTowerDataForClassParts(newContexts) {
+        return withNewTowerDataForClass(newContexts) {
             f()
         }
     }